bugfix: run precip func once, then index if 3d

gemini3d · Oct 27, 2023 · b8969dd · b8969dd
1 parent c083bb2
commit b8969dd
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/gemini3d/particles/__init__.py b/src/gemini3d/particles/__init__.py
@@ -53,7 +53,10 @@ def particles_BCs(cfg: dict[str, T.Any], xg: dict[str, T.Any]):
     Qtmp = Qfunc(pg, cfg["Qprecip"], cfg["Qprecip_background"])
 
     for i in range(i_on, i_off):
-        pg["Q"][i, :, :] = Qtmp[i, :, :]
+        if Qtmp.ndim == 3:
+            pg["Q"][i, :, :] = Qtmp[i, :, :]
+        else:
+            pg["Q"][i, :, :] = Qtmp
         pg["E0"][i, :, :] = cfg["E0precip"]
 
     assert np.isfinite(pg["Q"]).all(), "Q flux must be finite"