saving gravitational potential and residual momenta for both standard time stepping and fitted p3m time stepping

src/wip3m/convergence_cola_p3m_only_expl.py

@@ -246,6 +246,8 @@ if __name__ == "__main__":
                 p3m_fit_coeffs = P3M_FIT_COEFFS_DEFAULT_2NP
             elif L == Np:
                 p3m_fit_coeffs = P3M_FIT_COEFFS_DEFAULT_NP
+            elif 2 * L == Np:
+                p3m_fit_coeffs = P3M_FIT_COEFFS_DEFAULT_05NP
             else:
                 raise ValueError(
                     f"Invalid box size L={L}. Must be either 2*Np={2*Np} or Np={Np} when using the fitted P3M limiter."
@@ -428,11 +430,12 @@ if __name__ == "__main__":
                     TimeStepDistribution=TimeStepDistribution,
                     nsteps=nsteps,
                     ymin=3e-3,
-                    fac_hubble=fac_H_custom,
-                    fac_bend=fac_bend,
+                    fac_hubble=parameters["fac_H_custom"],
+                    fac_bend=parameters["fac_bend"],
+                    fac_p3m_fit=parameters["fac_p3m_fit"],
                     plot_bend=not use_p3m_fit,
-                    da_max_early=da_early,
-                    da_max_late=DEFAULT_DA_MAX_LATE_CUSTOM,
+                    da_max_early=parameters["da_max_early_custom"],
+                    da_max_late=parameters["da_max_late_custom"],
                     save_path=outdir + f"time_step_diagnostics_custom_{name}.pdf",
                 )
 
@@ -482,7 +485,9 @@ if __name__ == "__main__":
 
         mod_val_reference = mod_val_p3m1
         Pk_ref = Pk_P3M1
-        label_ref = f"P3M, mod_val={mod_val_reference}, ns={nsteps_dict['p3m1']}"
+        label_ref = (
+            rf"P3M, {name_scaled}={mod_val_reference}, $n_\textrm{{steps}}$={nsteps_dict['p3m1']}"
+        )
 
         logger.info(f"Plotting power spectra...")
         INDENT()
@@ -516,7 +521,7 @@ if __name__ == "__main__":
             (field_name, Pk) for field_name, Pk in fields_to_plot if field_name != ref
         ]
         for field_name, Pk in fields_to_plot:
-            label = f"{field_name}, mod_val={eval(f'mod_val_{field_name.lower()}')}, ns={nsteps_dict[field_name.lower()]}"
+            label = rf"P3M, {name_scaled}={eval(f'mod_val_{field_name.lower()}')}, $n_\textrm{{steps}}$={nsteps_dict[field_name.lower()]}"
             linestyle = "--"
             zorder = 2
             ax.plot(k, Pk / Pk_ref, label=label, linestyle=linestyle)

src/wip3m/params.py

@@ -146,9 +146,36 @@ DEFAULT_SUB_BEND1_COLA = 0.012
 DEFAULT_SUB_BEND2_COLA = 0.014
 DEFAULT_FAC_P3M_FIT = 1.0
 DEFAULT_DA_MAX_LATE_CUSTOM = 0.02
-P3M_FIT_COEFFS_DEFAULT_2NP = [-4.41928363, 0.24380666, 1.92034984, 1.21626721, 0.81694801, 0.20590599]  # fitted P3M limiter for L=2Np
-P3M_FIT_COEFFS_DEFAULT_NP = [-5.14695921, 0.35234941, 2.08397701, 1.2447759, 0.61695249, 0.11401126]  # fitted P3M limiter for L=Np
-# TODO: add polynomial coefficients for more L/Np values
+
+# V1:
+# P3M_FIT_COEFFS_DEFAULT_2NP = [-4.41928363, 0.24380666, 1.92034984, 1.21626721, 0.81694801, 0.20590599]  # fitted P3M limiter for L=2Np
+# P3M_FIT_COEFFS_DEFAULT_NP = [-5.14695921, 0.35234941, 2.08397701, 1.2447759, 0.61695249, 0.11401126]  # fitted P3M limiter for L=Np
+
+# V2:
+P3M_FIT_COEFFS_DEFAULT_2NP = [
+    -2.421323,
+    0.45965703,
+    2.52817976,
+    1.58540987,
+    0.81381042,
+    0.17553455,
+]  # fitted P3M limiter for L=2Np based on COLAr6/V3cola_Pf_L1024_N1024_Np512
+P3M_FIT_COEFFS_DEFAULT_NP = [
+    -3.19544528,
+    -0.01948223,
+    0.5180241,
+    -1.08372045,
+    -0.71808255,
+    -0.14370312,
+]  # fitted P3M limiter for L=Np based on COLAr6/V3cola_Pf_L512_N1024_Np512
+P3M_FIT_COEFFS_DEFAULT_05NP = [
+    -3.74659557,
+    0.39041298,
+    1.90821078,
+    1.10022755,
+    0.5597775,
+    0.09604354,
+]  # fitted P3M limiter for L=Np based on COLAr7/CV2_cola_Pf_L80_N160_Np160
 
 # pyright: reportUnsupportedDunderAll=false
 __all__ = [k for k in globals() if not k.startswith("_")]

src/wip3m/plot_utils.py

@@ -520,8 +520,18 @@ def plot_custom_timestepping_diagnostics(
             lab = "Bend-approx. P3M"
         ax.loglog(a, da_bend / a, marker="o", markersize=4, linestyle="-", label=lab)
     if fac_p3m_fit is not None:
+        # lab = rf"P3M fit ($\eta_\textrm{{Pf}}={fac_p3m_fit:.3f}$)"
+        # ax.loglog(a, da_p3m_fit / a, marker="o", markersize=4, linestyle="-", label=lab)
+        # bottom_pad += 0.06
+        # For da_p3m_fit, only plot the points corresponding to a > 1e-1:
+        amin = 1e-1
+        mask = a > amin
+        if not np.any(mask):
+            raise ValueError(f"No values of 'a' greater than {amin} found in {log_path_custom}.")
+        a_p3m_fit = a[mask]
+        da_p3m_fit = da_p3m_fit[mask]
         lab = rf"P3M fit ($\eta_\textrm{{Pf}}={fac_p3m_fit:.3f}$)"
-        ax.loglog(a, da_p3m_fit / a, marker="o", markersize=4, linestyle="-", label=lab)
+        ax.loglog(a_p3m_fit, da_p3m_fit / a_p3m_fit, marker="o", markersize=4, linestyle="-", label=lab)
         bottom_pad += 0.06
     # _, _, _, _, _, _, _, da_p3m_concept, _ = np.loadtxt(log_path, delimiter=",", unpack=True, skiprows=0)
     # ax.loglog(a, da_p3m_concept / a, marker="o", markersize=4, linestyle="-", label=r"$\Delta t\propto x_s/\sqrt{\langle v^2\rangle}$ (P3M, untuned)")
@@ -585,11 +595,12 @@ def fit_p3m(a, da_p3m, degree=3, amin=1e-1):
         raise ValueError(f"No values of 'a' greater than {amin} found.")
     log_a = np.log(a[mask])
     log_y = np.log(da_p3m[mask] / a[mask])
-    if degree == 3:
-        initial_guess = [-3.0, 0.2, 2.5, 0.6]
-    elif degree == 5:
-        initial_guess = [-3.5, -0.3, -0.5, -4.0, -3.0, -0.5]
-    else:
+    # if degree == 3:
+    #     initial_guess = [-3.0, 0.2, 2.5, 0.6]
+    # elif degree == 5:
+    #     initial_guess = [-3.5, 0.2, 1.5, 1.0, 0.5, 0.15]
+    # else:
+    #     initial_guess = np.ones(degree + 1)
     initial_guess = np.ones(degree + 1)
     coeffs, _ = curve_fit(log_poly_model, log_a, log_y, p0=initial_guess)
     return coeffs

src/wip3m/tools.py

@@ -322,6 +322,7 @@ def generate_sim_params(params_dict, ICs, workdir, outdir, file_ext=None, force=
         else:
             sub_bend1_diagnostic = params_dict.get("sub_bend1", DEFAULT_SUB_BEND1)
             sub_bend2_diagnostic = params_dict.get("sub_bend2", DEFAULT_SUB_BEND2)
+        p3m_fit_coeffs = params_dict.get("p3m_fit_coeffs", P3M_FIT_COEFFS_DEFAULT_NP)
         S = param_file(
             # Basic setup:
             Particles=Particles,
@@ -365,8 +366,24 @@ def generate_sim_params(params_dict, ICs, workdir, outdir, file_ext=None, force=
             sub_bend2=params_dict.get("sub_bend2", sub_bend2_diagnostic),
             fac_p3m_fit=params_dict.get("fac_p3m_fit", DEFAULT_FAC_P3M_FIT),
             use_p3m_fit=params_dict.get("use_p3m_fit", False),
-            da_max_early_custom=params_dict.get("da_max_early_custom", DEFAULT_DA_MAX_EARLY_CUSTOM),
+            p3m_fit_coeff0=p3m_fit_coeffs[0],
+            p3m_fit_coeff1=p3m_fit_coeffs[1],
+            p3m_fit_coeff2=p3m_fit_coeffs[2],
+            p3m_fit_coeff3=p3m_fit_coeffs[3],
+            p3m_fit_coeff4=p3m_fit_coeffs[4],
+            p3m_fit_coeff5=p3m_fit_coeffs[5],
+            da_max_early_custom=params_dict.get(
+                "da_max_early_custom", DEFAULT_DA_MAX_EARLY_CUSTOM
+            ),
             da_max_late_custom=params_dict.get("da_max_late_custom", DEFAULT_DA_MAX_LATE_CUSTOM),
+            WriteGravPot=params_dict.get("WriteGravPot", False),
+            OutputGravitationalPotentialBase=params_dict.get(
+                "OutputGravitationalPotentialBase", simpath + "gravpot_"
+            ),
+            WriteReferenceFrame=params_dict.get("WriteReferenceFrame", False),
+            OutputMomentaBase=params_dict.get(
+                "OutputMomentaBase", simpath + "momenta_"
+            ),
             #############################
             ## Cosmological parameters ##
             #############################

submit/actual_cola_p3m.sh

@@ -1,11 +1,11 @@
 #!/bin/bash
-#SBATCH --job-name=cola_Pf_L512_N512_Np256
-#SBATCH --output=/data70/hoellinger/WIP3M/cola_Pf_L512_N512_Np256/log.log
-#SBATCH --error=/data70/hoellinger/WIP3M/cola_Pf_L512_N512_Np256/err.err
+#SBATCH --job-name=Stockholm_colap3m_LNNp64
+#SBATCH --output=/data70/hoellinger/WIP3M/Stockholm_colap3m_LNNp64/log.log
+#SBATCH --error=/data70/hoellinger/WIP3M/Stockholm_colap3m_LNNp64/err.err
 #SBATCH --nodes=1                        # Number of nodes (value or min-max)
 #SBATCH --ntasks=128                     # The number of tasks (i.e. cores) per node
 #SBATCH --partition=comp,pscomp          # Partition name
-#SBATCH --time=48:00:00
+#SBATCH --time=24:00:00
 
 ##SBATCH --exclusive
 ##SBATCH --nodelist=i26             # Node name
@@ -15,60 +15,60 @@
 
 conda activate p3m
 
+export OMP_NUM_THREADS=32
+python $WIP3M_ROOT_PATH"src/wip3m/convergence_cola_p3m_only_expl.py" \
+    --run_id Stockholm_colap3m_LNNp64 \
+    --L 64 \
+    --N 64 \
+    --Np 64 \
+    --Npm 128 \
+    --n_Tiles 16 \
+    --z_i 19.0 \
+    --z_f 0.0 \
+    --plot_fields True \
+    --rescale_limiter "fac_p3m_fit" \
+    --use_p3m_fit True \
+    --mod_val_p3m1 0.03 \
+    --mod_val_p3m2 0.07 \
+    --mod_val_p3m3 0.15
+
 # export OMP_NUM_THREADS=64
 # python $WIP3M_ROOT_PATH"src/wip3m/convergence_cola_p3m_only_expl.py" \
-#     --run_id cola_Pf_LNNp512 \
-#     --L 512 \
-#     --N 512 \
-#     --Np 512 \
-#     --Npm 1024 \
-#     --n_Tiles 128 \
+#     --run_id Stockholm_colap3m_LNNp160 \
+#     --L 160 \
+#     --N 160 \
+#     --Np 160 \
+#     --Npm 320 \
+#     --n_Tiles 48 \
 #     --z_i 19.0 \
 #     --z_f 0.0 \
 #     --plot_fields True \
 #     --rescale_limiter "fac_p3m_fit" \
 #     --use_p3m_fit True \
-#     --mod_val_p3m1 0.3 \
-#     --mod_val_p3m2 0.5 \
-#     --mod_val_p3m3 1.1
+#     --mod_val_p3m1 0.03 \
+#     --mod_val_p3m2 0.07 \
+#     --mod_val_p3m3 0.15
 
 # export OMP_NUM_THREADS=64
 # python $WIP3M_ROOT_PATH"src/wip3m/convergence_cola_p3m_only_expl.py" \
-#     --run_id cola_Pf_L1024_N512_Np512 \
-#     --L 1024 \
-#     --N 512 \
-#     --Np 512 \
-#     --Npm 1024 \
-#     --n_Tiles 128 \
+#     --run_id Stockholm_colap3m_L80_N80_Np160 \
+#     --L 80 \
+#     --N 80 \
+#     --Np 160 \
+#     --Npm 320 \
+#     --n_Tiles 48 \
 #     --z_i 19.0 \
 #     --z_f 0.0 \
 #     --plot_fields True \
 #     --rescale_limiter "fac_p3m_fit" \
 #     --use_p3m_fit True \
-#     --mod_val_p3m1 0.33 \
-#     --mod_val_p3m2 0.4 \
-#     --mod_val_p3m3 0.5
+#     --mod_val_p3m1 0.05 \
+#     --mod_val_p3m2 0.07 \
+#     --mod_val_p3m3 0.15
 
-# export OMP_NUM_THREADS=64
+# export OMP_NUM_THREADS=65
 # python $WIP3M_ROOT_PATH"src/wip3m/convergence_cola_p3m_only_expl.py" \
-#     --run_id cola_Pf_L1024_N1024_Np512 \
-#     --L 1024 \
-#     --N 1024 \
-#     --Np 512 \
-#     --Npm 1024 \
-#     --n_Tiles 128 \
-#     --z_i 19.0 \
-#     --z_f 0.0 \
-#     --plot_fields True \
-#     --rescale_limiter "fac_p3m_fit" \
-#     --use_p3m_fit True \
-#     --mod_val_p3m1 0.33 \
-#     --mod_val_p3m2 0.4 \
-#     --mod_val_p3m3 0.5
-
-# export OMP_NUM_THREADS=64
-# python $WIP3M_ROOT_PATH"src/wip3m/convergence_cola_p3m_only_expl.py" \
-#     --run_id cola_Pf_LNNp256 \
+#     --run_id Stockholm_colap3m_LNNp256 \
 #     --L 256 \
 #     --N 256 \
 #     --Np 256 \
@@ -79,15 +79,15 @@ conda activate p3m
 #     --plot_fields True \
 #     --rescale_limiter "fac_p3m_fit" \
 #     --use_p3m_fit True \
-#     --mod_val_p3m1 0.3 \
-#     --mod_val_p3m2 0.5 \
-#     --mod_val_p3m3 1.1
+#     --mod_val_p3m1 0.03 \
+#     --mod_val_p3m2 0.07 \
+#     --mod_val_p3m3 0.15
 
-# export OMP_NUM_THREADS=64
+# export OMP_NUM_THREADS=65
 # python $WIP3M_ROOT_PATH"src/wip3m/convergence_cola_p3m_only_expl.py" \
-#     --run_id cola_Pf_L512_N256_Np256 \
-#     --L 512 \
-#     --N 256 \
+#     --run_id Stockholm_colap3m_L128_N128_Np256 \
+#     --L 128 \
+#     --N 128 \
 #     --Np 256 \
 #     --Npm 512 \
 #     --n_Tiles 64 \
@@ -96,25 +96,6 @@ conda activate p3m
 #     --plot_fields True \
 #     --rescale_limiter "fac_p3m_fit" \
 #     --use_p3m_fit True \
-#     --mod_val_p3m1 0.33 \
-#     --mod_val_p3m2 0.4 \
-#     --mod_val_p3m3 0.5
-
-export OMP_NUM_THREADS=64
-python $WIP3M_ROOT_PATH"src/wip3m/convergence_cola_p3m_only_expl.py" \
-    --run_id cola_Pf_L512_N512_Np256 \
-    --L 512 \
-    --N 512 \
-    --Np 256 \
-    --Npm 512 \
-    --n_Tiles 64 \
-    --z_i 19.0 \
-    --z_f 0.0 \
-    --plot_fields True \
-    --rescale_limiter "fac_p3m_fit" \
-    --use_p3m_fit True \
-    --mod_val_p3m1 0.33 \
-    --mod_val_p3m2 0.4 \
-    --mod_val_p3m3 0.5
-
-exit 0
+#     --mod_val_p3m1 0.03 \
+#     --mod_val_p3m2 0.07 \
+#     --mod_val_p3m3 0.15