Merge PDE experiment code into main

src/gen_experiments/config.py

@@ -102,10 +102,10 @@ def addn(x):
     "med-noise": ND({"n_trajectories": 2, "noise_abs": 0.8}),
     "med-noise-many": ND({"n_trajectories": 10, "noise_abs": 0.8}),
     "hi-noise": ND({"n_trajectories": 2, "noise_abs": 2}),
-    "pde-ic1": ND({"init_cond": np.exp(-((np.arange(0, 10, 0.1) - 5) ** 2) / 2)}),
+    "pde-ic1": ND({"init_cond": np.exp(-((np.linspace(-8, 8, 256) + 2) ** 2) / 2)}),
     "pde-ic2": ND({
-        "init_cond": (np.cos(np.arange(0, 10, 0.1))) * (
-            1 + np.sin(np.arange(0, 10, 0.1) - 0.5)
+        "init_cond": (np.cos(np.linspace(0, 100, 1024))) * (
+            1 + np.sin(np.linspace(0, 100, 1024) - 0.5)
         )
     }),
 }
@@ -131,36 +131,23 @@ def addn(x):
     "testweak": ND({"featcls": "WeakPDELibrary"}),  # needs work
     "pde2": ND({
         "featcls": "pde",
-        "library_functions": [identity, quadratic],
-        "function_names": [identity, addn],
+        "function_library": ps.PolynomialLibrary(degree=2, include_bias=False),
         "derivative_order": 2,
-        "spatial_grid": np.arange(0, 10, 0.1),
+        "spatial_grid": np.linspace(-8, 8, 256),
         "include_interaction": True,
     }),
-    "pde3": ND({
-        "featcls": "pde",
-        "library_functions": [identity, quadratic],
-        "function_names": [identity, addn],
-        "derivative_order": 3,
-        "spatial_grid": np.arange(0, 10, 0.1),
-        "include_interaction": True,
-        "is_uniform": True,
-    }),
     "pde4": ND({
         "featcls": "pde",
-        "library_functions": [identity, quadratic],
-        "function_names": [identity, addn],
+        "function_library": ps.PolynomialLibrary(degree=2, include_bias=False),
         "derivative_order": 4,
-        "spatial_grid": np.arange(0, 10, 0.1),
+        "spatial_grid": np.linspace(0, 100, 1024),
         "include_interaction": True,
-        "is_uniform": True,
-        "periodic": True,
-        "include_bias": True,
     }),
 }
 opt_params = {
     "test": ND({"optcls": "STLSQ"}),
-    "miosr": ND({"optcls": "MIOSR", "unbias": True}),
+    "test_low": ND({"optcls": "STLSQ", "threshold": 0.09}),
+    "miosr": ND({"optcls": "MIOSR"}),
     "enslsq": ND(
         {"optcls": "ensemble", "opt": ps.STLSQ(), "bagging": True, "n_models": 20}
     ),
@@ -200,6 +187,18 @@ def addn(x):
         "feat_params": feat_params["test"],
         "opt_params": opt_params["test"],
     }),
+    "test-pde1": ND({
+        "sim_params": sim_params["pde-ic1"],
+        "diff_params": diff_params["test_axis"],
+        "feat_params": feat_params["pde2"],
+        "opt_params": opt_params["test_low"],
+    }),
+    "test-pde2": ND({
+        "sim_params": sim_params["pde-ic2"],
+        "diff_params": diff_params["test_axis"],
+        "feat_params": feat_params["pde4"],
+        "opt_params": opt_params["test"],
+    }),
     "tv1": ND({
         "sim_params": sim_params["test"],
         "diff_params": diff_params["tv"],

src/gen_experiments/pdes.py

@@ -54,7 +54,7 @@ def burgers1D_periodic(t, u, dx, nx):
     u = np.reshape(u, nx)
     uxx = SpectralDerivative(d=2, axis=0)._differentiate(u, dx)
     ux = SpectralDerivative(d=1, axis=0)._differentiate(u, dx)
-    return np.reshape((uxx - u * ux), nx)
+    return np.reshape((0.1 * uxx - u * ux), nx)
 
 
 def ks_dirichlet(t, u, dx, nx):
@@ -88,39 +88,34 @@ def kdv(t, u, dx, nx):
     "diffuse1D_dirichlet": {
         "rhsfunc": {"func": diffuse1D_dirichlet, "dimension": 1},
         "input_features": ["u"],
-        "spatial_args": [0.1, 100],
         "time_args": [0.1, 10],
         "coeff_true": [{"u_11": 1}],
         "spatial_grid": np.arange(0, 10, 0.1),
     },
     "diffuse1D_periodic": {
         "rhsfunc": {"func": diffuse1D_periodic, "dimension": 1},
         "input_features": ["u"],
-        "spatial_args": [0.1, 100],
         "time_args": [0.1, 10],
         "coeff_true": [{"u_11": 1}],
-        "spatial_grid": np.arange(0, 10, 0.1),
+        "spatial_grid": np.linspace(-8, 8, 256),
     },
     "burgers1D_dirichlet": {
         "rhsfunc": {"func": burgers1D_dirichlet, "dimension": 1},
         "input_features": ["u"],
-        "spatial_args": [0.1, 100],
         "time_args": [0.1, 10],
         "coeff_true": [{"u_11": 1, "uu_1": 1}],
         "spatial_grid": np.arange(0, 10, 0.1),
     },
     "burgers1D_periodic": {
         "rhsfunc": {"func": burgers1D_periodic, "dimension": 1},
         "input_features": ["u"],
-        "spatial_args": [0.1, 100],
         "time_args": [0.1, 10],
-        "coeff_true": [{"u_11": 1, "uu_1": -1}],
-        "spatial_grid": np.arange(0, 10, 0.1),
+        "coeff_true": [{"u_11": 0.1, "uu_1": -1}],
+        "spatial_grid": np.linspace(-8, 8, 256),
     },
     "ks_dirichlet": {
         "rhsfunc": {"func": ks_dirichlet, "dimension": 1},
         "input_features": ["u"],
-        "spatial_args": [0.1, 100],
         "time_args": [0.1, 10],
         "coeff_true": [
             {"u_11": -1, "u_1111": -1, "uu_1": -1},
@@ -130,12 +125,11 @@ def kdv(t, u, dx, nx):
     "ks_periodic": {
         "rhsfunc": {"func": ks_periodic, "dimension": 1},
         "input_features": ["u"],
-        "spatial_args": [0.1, 100],
-        "time_args": [0.1, 10],
+        "time_args": [0.4, 100],
         "coeff_true": [
             {"u_11": -1, "u_1111": -1, "uu_1": -1},
         ],
-        "spatial_grid": np.arange(0, 10, 0.1),
+        "spatial_grid": np.linspace(0, 100, 1024),
     },
 }
 
@@ -153,7 +147,11 @@ def run(
     rhsfunc = pde_setup[group]["rhsfunc"]["func"]
     input_features = pde_setup[group]["input_features"]
     initial_condition = sim_params["init_cond"]
-    spatial_args = pde_setup[group]["spatial_args"]
+    spatial_grid = pde_setup[group]["spatial_grid"]
+    spatial_args = [
+        (spatial_grid[-1] - spatial_grid[0]) / len(spatial_grid),
+        len(spatial_grid),
+    ]
     time_args = pde_setup[group]["time_args"]
     dimension = pde_setup[group]["rhsfunc"]["dimension"]
     coeff_true = pde_setup[group]["coeff_true"]
@@ -197,7 +195,11 @@ def run(
         )
         trial_data["model"].print()
         plot_pde_training_data(
-            trial_data["x_train"], trial_data["x_true"], trial_data["smooth_train"]
+            spatial_grid,
+            trial_data["t_train"],
+            trial_data["x_train"],
+            trial_data["x_true"],
+            trial_data["smooth_train"],
         )
         compare_coefficient_plots(
             trial_data["coeff_fit"],