ENH: PDE initial conditions assigned to sim_params

Jacob-Stevens-Haas · Jan 30, 2024 · 9169b5c · 9169b5c
1 parent 324d3a9
commit 9169b5c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 12 deletions.
diff --git a/src/gen_experiments/config.py b/src/gen_experiments/config.py
@@ -130,6 +130,13 @@ 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-ic2": ND({"init_cond": np.exp(-((np.arange(0, 10, 0.1) - 5) ** 2) / 2)}),
+    "pde-ic3": ND({
+        "init_cond": (np.cos(np.arange(0, 10, 0.1))) * (
+            1 + np.sin(np.arange(0, 10, 0.1) - 0.5)
+        )
+    }),
 }
 diff_params = {
     "test": ND({"diffcls": "FiniteDifference"}),

diff --git a/src/gen_experiments/pdes.py b/src/gen_experiments/pdes.py
@@ -89,7 +89,6 @@ def kdv(t, u, dx, nx):
     "diffuse1D_dirichlet": {
         "rhsfunc": {"func": diffuse1D_dirichlet, "dimension": 1},
         "input_features": ["u"],
-        "initial_condition": 10 * np.exp(-((np.arange(0, 10, 0.1) - 5) ** 2) / 2),
         "spatial_args": [0.1, 100],
         "time_args": [0.1, 10],
         "coeff_true": [{"u_11": 1}],
@@ -98,7 +97,6 @@ def kdv(t, u, dx, nx):
     "diffuse1D_periodic": {
         "rhsfunc": {"func": diffuse1D_periodic, "dimension": 1},
         "input_features": ["u"],
-        "initial_condition": np.exp(-((np.arange(0, 10, 0.1) - 5) ** 2) / 2),
         "spatial_args": [0.1, 100],
         "time_args": [0.1, 10],
         "coeff_true": [{"u_11": 1}],
@@ -107,7 +105,6 @@ def kdv(t, u, dx, nx):
     "burgers1D_dirichlet": {
         "rhsfunc": {"func": burgers1D_dirichlet, "dimension": 1},
         "input_features": ["u"],
-        "initial_condition": 10 * np.exp(-((np.arange(0, 10, 0.1) - 5) ** 2) / 2),
         "spatial_args": [0.1, 100],
         "time_args": [0.1, 10],
         "coeff_true": [{"u_11": 1, "uu_1": 1}],
@@ -116,7 +113,6 @@ def kdv(t, u, dx, nx):
     "burgers1D_periodic": {
         "rhsfunc": {"func": burgers1D_periodic, "dimension": 1},
         "input_features": ["u"],
-        "initial_condition": np.exp(-((np.arange(0, 10, 0.1) - 5) ** 2) / 2),
         "spatial_args": [0.1, 100],
         "time_args": [0.1, 10],
         "coeff_true": [{"u_11": 1, "uu_1": -1}],
@@ -125,9 +121,6 @@ def kdv(t, u, dx, nx):
     "ks_dirichlet": {
         "rhsfunc": {"func": ks_dirichlet, "dimension": 1},
         "input_features": ["u"],
-        "initial_condition": np.cos(np.arange(0, 10, 0.1)) * (
-            1 + np.sin(np.arange(0, 10, 0.1))
-        ),
         "spatial_args": [0.1, 100],
         "time_args": [0.1, 10],
         "coeff_true": [
@@ -138,9 +131,6 @@ def kdv(t, u, dx, nx):
     "ks_periodic": {
         "rhsfunc": {"func": ks_periodic, "dimension": 1},
         "input_features": ["u"],
-        "initial_condition": (np.cos(np.arange(0, 10, 0.1))) * (
-            1 + np.sin(np.arange(0, 10, 0.1) - 0.5)
-        ),
         "spatial_args": [0.1, 100],
         "time_args": [0.1, 10],
         "coeff_true": [
@@ -154,6 +144,7 @@ def kdv(t, u, dx, nx):
 def run(
     seed: float,
     group: str,
+    sim_params: dict,
     diff_params: dict,
     feat_params: dict,
     opt_params: dict,
@@ -162,7 +153,6 @@ def run(
 ) -> dict | tuple[dict, TrialData | FullTrialData]:
     rhsfunc = pde_setup[group]["rhsfunc"]["func"]
     input_features = pde_setup[group]["input_features"]
-    initial_condition = pde_setup[group]["initial_condition"]
     spatial_args = pde_setup[group]["spatial_args"]
     time_args = pde_setup[group]["time_args"]
     dimension = pde_setup[group]["rhsfunc"]["dimension"]
@@ -173,7 +163,7 @@ def run(
         time_args = [0.01, 10]
     dt, t_train, x_train, x_test, x_dot_test, x_train_true = gen_pde_data(
         rhsfunc,
-        initial_condition,
+        sim_params["init_cond"],
         spatial_args,
         dimension,
         seed,