black args fix

mikegrudic · Sep 25, 2024 · d36e8bf · d36e8bf
1 parent 467ff56
commit d36e8bf
Show file tree

Hide file tree

Showing 3 changed files with 393 additions and 235 deletions.
diff --git a/MakeCloud.py b/MakeCloud.py
@@ -116,16 +116,24 @@ def TurbField(res=256, minmode=2, maxmode=64, sol_weight=1.0, seed=42):
             if i == j:
                 vk_new[i] += sol_weight * VK[j]
             vk_new[i] += (
-                (1 - 2 * sol_weight) * freq3d[i] * freq3d[j] / (kSqr + 1e-300) * VK[j]
+                (1 - 2 * sol_weight)
+                * freq3d[i]
+                * freq3d[j]
+                / (kSqr + 1e-300)
+                * VK[j]
             )
     vk_new[:, kSqr == 0] = 0.0
     VK = vk_new
 
     vel = np.array(
         [fftpack.ifftn(vk).real for vk in VK]
     )  # transform back to real space
-    vel -= np.average(vel, axis=(1, 2, 3))[:, np.newaxis, np.newaxis, np.newaxis]
-    vel = vel / np.sqrt(np.sum(vel**2, axis=0).mean())  # normalize so that RMS is 1
+    vel -= np.average(vel, axis=(1, 2, 3))[
+        :, np.newaxis, np.newaxis, np.newaxis
+    ]
+    vel = vel / np.sqrt(
+        np.sum(vel**2, axis=0).mean()
+    )  # normalize so that RMS is 1
     return np.array(vel)
 
 
@@ -179,8 +187,9 @@ def TurbField(res=256, minmode=2, maxmode=64, sol_weight=1.0, seed=42):
 
         print("Downloading glass file...")
         urllib.request.urlretrieve(
-            "http://www.tapir.caltech.edu/~mgrudich/glass_orig.npy", glass_path
-#            "https://data.obs.carnegiescience.edu/starforge/glass_orig.npy", glass_path
+            "http://www.tapir.caltech.edu/~mgrudich/glass_orig.npy",
+            glass_path,
+            #            "https://data.obs.carnegiescience.edu/starforge/glass_orig.npy", glass_path
         )
 
 if localdir:
@@ -220,7 +229,8 @@ def TurbField(res=256, minmode=2, maxmode=64, sol_weight=1.0, seed=42):
     )
     + ("_%d" % seed)
     + (
-        "_collision_%g_%g_%g_%s" % (impact_dist, impact_param, v_impact, impact_axis)
+        "_collision_%g_%g_%g_%s"
+        % (impact_dist, impact_param, v_impact, impact_axis)
         if impact_dist > 0
         else ""
     )
@@ -233,7 +243,9 @@ def TurbField(res=256, minmode=2, maxmode=64, sol_weight=1.0, seed=42):
 delta_m_solar = delta_m / mass_unit
 rho_avg = 3 * M_gas / R**3 / (4 * np.pi)
 if delta_m_solar < 0.1:  # if we're doing something marginally IMF-resolving
-    softening = 3.11e-5  # ~6.5 AU, minimum sink radius is 2.8 times that (~18 AU)
+    softening = (
+        3.11e-5  # ~6.5 AU, minimum sink radius is 2.8 times that (~18 AU)
+    )
     ncrit = 1e13  # ~100x the opacity limit
 else:  # something more FIRE-like, where we rely on a sub-grid prescription turning gas into star particles
     softening = 0.1
@@ -256,7 +268,9 @@ def TurbField(res=256, minmode=2, maxmode=64, sol_weight=1.0, seed=42):
 )  # ST_Energy sets the dissipation rate of SPECIFIC energy ~ v^2 / (L/v) ~ v^3/L
 
 paramsfile = str(
-    open(os.path.realpath(__file__).replace("MakeCloud.py", "params.txt"), "r").read()
+    open(
+        os.path.realpath(__file__).replace("MakeCloud.py", "params.txt"), "r"
+    ).read()
 )
 
 jet_particle_mass = min(delta_m, max(1e-4, delta_m / 10.0))
@@ -295,7 +309,9 @@ def TurbField(res=256, minmode=2, maxmode=64, sol_weight=1.0, seed=42):
 }
 
 for k, r in replacements.items():
-    paramsfile = paramsfile.replace(k,(r if isinstance(r,str) else '{:.2e}'.format(r)))
+    paramsfile = paramsfile.replace(
+        k, (r if isinstance(r, str) else "{:.2e}".format(r))
+    )
 
 open("params_" + filename.replace(".hdf5", "") + ".txt", "w").write(paramsfile)
 if makebox:
@@ -306,12 +322,15 @@ def TurbField(res=256, minmode=2, maxmode=64, sol_weight=1.0, seed=42):
     replacements_box["TURB_MAXLAMBDA"] = int(100 * L * 2) / 100
     paramsfile = str(
         open(
-            os.path.realpath(__file__).replace("MakeCloud.py", "params.txt"), "r"
+            os.path.realpath(__file__).replace("MakeCloud.py", "params.txt"),
+            "r",
         ).read()
     )
     for k in replacements_box.keys():
         paramsfile = paramsfile.replace(k, str(replacements_box[k]))
-    open("params_" + filename.replace(".hdf5", "") + "_BOX.txt", "w").write(paramsfile)
+    open("params_" + filename.replace(".hdf5", "") + "_BOX.txt", "w").write(
+        paramsfile
+    )
 if makecylinder:
     # Get cylinder params
     R_cyl = R * np.sqrt(
@@ -346,12 +365,15 @@ def TurbField(res=256, minmode=2, maxmode=64, sol_weight=1.0, seed=42):
     replacements_cyl["TURB_KMAX"] = int(100 * 4 * np.pi / (R_cyl) + 1) / 100.0
     paramsfile = str(
         open(
-            os.path.realpath(__file__).replace("MakeCloud.py", "params.txt"), "r"
+            os.path.realpath(__file__).replace("MakeCloud.py", "params.txt"),
+            "r",
         ).read()
     )
     for k in replacements_cyl.keys():
         paramsfile = paramsfile.replace(k, str(replacements_cyl[k]))
-    open("params_" + filename.replace(".hdf5", "") + "_CYL.txt", "w").write(paramsfile)
+    open("params_" + filename.replace(".hdf5", "") + "_CYL.txt", "w").write(
+        paramsfile
+    )
 
 if param_only:
     print("Parameters only run, exiting...")
@@ -414,9 +436,16 @@ def TurbField(res=256, minmode=2, maxmode=64, sol_weight=1.0, seed=42):
 
 B = np.c_[np.zeros(N_gas), np.zeros(N_gas), np.ones(N_gas)]
 vA_unit = (
-    3.429e8 * B_unit * (M_gas) ** -0.5 * R**1.5 * np.sqrt(4 * np.pi / 3) / v_unit
+    3.429e8
+    * B_unit
+    * (M_gas) ** -0.5
+    * R**1.5
+    * np.sqrt(4 * np.pi / 3)
+    / v_unit
 )  # alfven speed for unit magnetic field
-uB = 0.5 * M_gas * vA_unit**2  # magnetic energy we would have for unit magnetic field
+uB = (
+    0.5 * M_gas * vA_unit**2
+)  # magnetic energy we would have for unit magnetic field
 if bfixed > 0:
     B = B * bfixed
 else:
@@ -432,7 +461,9 @@ def TurbField(res=256, minmode=2, maxmode=64, sol_weight=1.0, seed=42):
 phi += phimode * np.sin(2 * phi) / 2
 x = (
     r[:, np.newaxis]
-    * np.c_[np.cos(phi) * np.sin(theta), np.sin(phi) * np.sin(theta), np.cos(theta)]
+    * np.c_[
+        np.cos(phi) * np.sin(theta), np.sin(phi) * np.sin(theta), np.cos(theta)
+    ]
 )
 
 if makecylinder:
@@ -519,13 +550,17 @@ def ind_in_cylinder(x, L_cyl, R_cyl):
         x_warm = x_warm[np.max(np.abs(x_warm), axis=1) < boxsize / 2]
         N_warm = len(x_warm)
         R_warm = (x_warm * x_warm).sum(1) ** 0.5
-        mgas = np.concatenate([mgas, np.clip(dm * (R_warm / (3 * R)) ** 3, dm, np.inf)])
+        mgas = np.concatenate(
+            [mgas, np.clip(dm * (R_warm / (3 * R)) ** 3, dm, np.inf)]
+        )
     else:
         x_warm = boxsize * np.random.rand(N_warm, 3) - boxsize / 2
         if impact_dist == 0:
             x_warm = x_warm[np.sum(x_warm**2, axis=1) > R**2]
         N_warm = len(x_warm)
-        mgas = np.concatenate([mgas, np.repeat(mgas.sum() / len(mgas), N_warm)])
+        mgas = np.concatenate(
+            [mgas, np.repeat(mgas.sum() / len(mgas), N_warm)]
+        )
     x = np.concatenate([x, x_warm])
     v = np.concatenate([v, np.zeros((N_warm, 3))])
     Bmag = np.average(np.sum(B**2, axis=1)) ** 0.5
@@ -578,7 +613,14 @@ def ind_in_cylinder(x, L_cyl, R_cyl):
     0,
     (1 if M_star > 0 else 0),
 ]
-F["Header"].attrs["NumPart_Total"] = [len(mgas), 0, 0, 0, 0, (1 if M_star > 0 else 0)]
+F["Header"].attrs["NumPart_Total"] = [
+    len(mgas),
+    0,
+    0,
+    0,
+    0,
+    (1 if M_star > 0 else 0),
+]
 F["Header"].attrs["BoxSize"] = boxsize
 F["Header"].attrs["Time"] = 0.0
 F["PartType0"].create_dataset("Masses", data=mgas)
@@ -591,7 +633,9 @@ def ind_in_cylinder(x, L_cyl, R_cyl):
     F.create_group("PartType5")
     # Let's add the sink at the center
     F["PartType5"].create_dataset("Masses", data=np.array([M_star]))
-    F["PartType5"].create_dataset("Coordinates", data=[x_star])  # at the center
+    F["PartType5"].create_dataset(
+        "Coordinates", data=[x_star]
+    )  # at the center
     F["PartType5"].create_dataset("Velocities", data=[v_star])  # at rest
     F["PartType5"].create_dataset(
         "ParticleIDs", data=np.array([F["PartType0/ParticleIDs"][:].max() + 1])
@@ -633,7 +677,9 @@ def ind_in_cylinder(x, L_cyl, R_cyl):
     F["PartType5"].create_dataset(
         "ProtoStellarRadius_inSolar", data=np.array([R_ZAMS])
     )  # Sinkradius set to softening
-    F["PartType5"].create_dataset("StarLuminosity_Solar", data=np.array([0.0]))  # dummy
+    F["PartType5"].create_dataset(
+        "StarLuminosity_Solar", data=np.array([0.0])
+    )  # dummy
     F["PartType5"].create_dataset("Mass_D", data=np.array([0.0]))  # No D left
 
 if magnetic_field > 0.0:
@@ -651,7 +697,8 @@ def ind_in_cylinder(x, L_cyl, R_cyl):
     F["Header"].attrs["Time"] = 0.0
     F["PartType0"].create_dataset("Masses", data=mgas[: len(mgas)])
     F["PartType0"].create_dataset(
-        "Coordinates", data=np.random.rand(len(mgas), 3) * F["Header"].attrs["BoxSize"]
+        "Coordinates",
+        data=np.random.rand(len(mgas), 3) * F["Header"].attrs["BoxSize"],
     )
     F["PartType0"].create_dataset("Velocities", data=np.zeros((len(mgas), 3)))
     F["PartType0"].create_dataset("ParticleIDs", data=1 + np.arange(len(mgas)))
@@ -672,7 +719,9 @@ def ind_in_cylinder(x, L_cyl, R_cyl):
     F["PartType0"].create_dataset("Masses", data=mgas)
     F["PartType0"].create_dataset("Coordinates", data=x_cyl)
     F["PartType0"].create_dataset("Velocities", data=v_cyl)
-    F["PartType0"].create_dataset("ParticleIDs", data=1 + np.arange(N_gas + N_warm_cyl))
+    F["PartType0"].create_dataset(
+        "ParticleIDs", data=1 + np.arange(N_gas + N_warm_cyl)
+    )
     F["PartType0"].create_dataset("InternalEnergy", data=u)
     if magnetic_field > 0.0:
         F["PartType0"].create_dataset("MagneticField", data=B_cyl)