magcalc: fix 2d/3d bugs as in MatGemini

src/gemini3d/hdf5/write.py

@@ -17,7 +17,7 @@
 CLVL = 3  # GZIP compression level: larger => better compression, slower to write
 
 
-def state(fn: Path, dat: xarray.Dataset):
+def state(fn: Path, dat: xarray.Dataset) -> None:
     """
     write STATE VARIABLE initial conditions
 
@@ -42,7 +42,7 @@ def state(fn: Path, dat: xarray.Dataset):
             _write_var(f, "/Phiall", dat["Phitop"])
 
 
-def _write_var(fid, name: str, A: xarray.DataArray):
+def _write_var(fid, name: str, A: xarray.DataArray) -> None:
     """
     NOTE: The .transpose() reverses the dimension order.
     The HDF Group never implemented the intended H5T_array_create(..., perm)
@@ -81,7 +81,7 @@ def _write_var(fid, name: str, A: xarray.DataArray):
     )
 
 
-def grid(size_fn: Path, grid_fn: Path, xg: dict[str, T.Any]):
+def grid(size_fn: Path, grid_fn: Path, xg: dict[str, T.Any]) -> None:
     """writes grid to disk
 
     Parameters
@@ -303,7 +303,7 @@ def neutral(fn: Path, N) -> None:
             )
 
 
-def maggrid(fn: Path, mag: dict[str, T.Any], gridsize: tuple[int, int, int]):
+def maggrid(fn: Path, mag: dict[str, T.Any], gridsize: tuple[int, int, int]) -> None:
     """
     hdf5 files can optionally store a gridsize variable which tells readers how to
     reshape the data into 2D or 3D arrays.
@@ -322,7 +322,7 @@ def maggrid(fn: Path, mag: dict[str, T.Any], gridsize: tuple[int, int, int]):
         f["/gridsize"] = np.asarray(gridsize).astype(np.int32)
 
 
-def write_time(fid: h5py.File, time: datetime):
+def write_time(fid: h5py.File, time: datetime) -> None:
     """
     write time to HDF5 file as year-month-day, UTsec
     """

src/gemini3d/magcalc.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 import argparse
 from pathlib import Path
-import logging
 from math import pi, radians
 import typing as T
 
@@ -11,7 +10,13 @@
 from . import write
 
 
-def magcalc(direc: Path, dang: float, xg: dict[str, T.Any] | None = None):
+def magcalc(
+    direc: Path,
+    dang: float,
+    Ltheta: int = 40,
+    Lphi: int = 40,
+    xg: dict[str, T.Any] | None = None,
+):
     """
     Parameters
     ----------
@@ -20,29 +25,26 @@ def magcalc(direc: Path, dang: float, xg: dict[str, T.Any] | None = None):
       top-level simulation path
 
     dang: float, optional
-      ANGULAR RANGE TO COVER FOR THE CALCLUATIONS.
-      THIS IS FOR THE FIELD POINTS - SOURCE POINTS COVER ENTIRE GRID.
+      ANGULAR RANGE TO COVER FOR THE FIELD POINTS - SOURCE POINTS COVER ENTIRE GRID
+
+    Ltheta: int, optional
+        number of points in theta
+
+    Lphi: int, optional
+        number of points in phi
 
     xg: dict, optional
       simulation grid
     """
 
     direc = Path(direc).expanduser()
-    if not direc.is_dir():
-        raise NotADirectoryError(direc)
 
     cfg = read.config(direc)
 
     if not xg:
         xg = read.grid(direc)
-        lx3 = xg["lx"][2]
 
-    if lx3 == 1:
-        flag2D = True
-        logging.info("2D meshgrid")
-    else:
-        flag2D = False
-        logging.info("3D meshgrid")
+    flag2D = any(xg["lx"][1:] == 1)
 
     # %% TABULATE THE SOURCE OR GRID CENTER LOCATION
     if "sourcemlon" not in cfg:
@@ -52,19 +54,20 @@ def magcalc(direc: Path, dang: float, xg: dict[str, T.Any] | None = None):
         thdist = pi / 2 - radians(cfg["sourcemlat"])  # zenith angle of source location
         phidist = radians(cfg["sourcemlon"])
 
-    # %% FIELD POINTS OF INTEREST (CAN/SHOULD BE DEFINED INDEPENDENT OF SIMULATION GRID)
-    ltheta = 10
-    lphi = 1 if flag2D else 10
+    # %% FIELD POINTS OF INTEREST
+    # CAN/SHOULD BE DEFINED INDEPENDENT OF SIMULATION GRID
+    if flag2D:
+        Lphi = 1
 
     thmin = thdist - radians(dang)
     thmax = thdist + radians(dang)
     phimin = phidist - radians(dang)
     phimax = phidist + radians(dang)
 
-    theta = np.linspace(thmin, thmax, ltheta)
-    phi = phidist if flag2D else np.linspace(phimin, phimax, lphi)
+    theta = np.linspace(thmin, thmax, Ltheta)
+    phi = phidist if flag2D else np.linspace(phimin, phimax, Lphi)
 
-    r = 6370e3 * np.ones((ltheta, lphi))
+    r = 6370e3 * np.ones((Ltheta, Lphi))
     # use ground level for altitude for all field points
     phi, theta = np.meshgrid(phi, theta)
 
@@ -74,8 +77,8 @@ def magcalc(direc: Path, dang: float, xg: dict[str, T.Any] | None = None):
         "phi": phi,
         "theta": theta,
     }
-    filename = direc / "inputs/magfieldpoints.h5"
-    write.maggrid(filename, mag)
+
+    write.maggrid(direc / "inputs/magfieldpoints.h5", mag)
 
 
 if __name__ == "__main__":

src/gemini3d/write.py

@@ -11,7 +11,7 @@
 from .hdf5 import write as h5write
 
 
-def state(out_file: Path, dat: xarray.Dataset, **kwargs):
+def state(out_file: Path, dat: xarray.Dataset, **kwargs) -> None:
     """
     WRITE STATE VARIABLE DATA.
     NOTE: WE don't write ANY OF THE ELECTRODYNAMIC
@@ -36,7 +36,7 @@ def state(out_file: Path, dat: xarray.Dataset, **kwargs):
     h5write.state(out_file, dat)
 
 
-def grid(cfg: dict[str, T.Any], xg: dict[str, T.Any]):
+def grid(cfg: dict[str, T.Any], xg: dict[str, T.Any]) -> None:
     """writes grid to disk
 
     Parameters
@@ -59,7 +59,7 @@ def grid(cfg: dict[str, T.Any], xg: dict[str, T.Any]):
     meta(input_dir / "setup_grid.json", git_meta(), cfg)
 
 
-def Efield(E: xarray.Dataset, outdir: Path):
+def Efield(E: xarray.Dataset, outdir: Path) -> None:
     """writes E-field to disk
 
     Parameters
@@ -77,7 +77,7 @@ def Efield(E: xarray.Dataset, outdir: Path):
     h5write.Efield(outdir, E)
 
 
-def precip(precip: xarray.Dataset, outdir: Path):
+def precip(precip: xarray.Dataset, outdir: Path) -> None:
     """writes precipitation to disk
 
     Parameters
@@ -109,7 +109,7 @@ def neutral2(data: dict[str, T.Any], outfile: Path):
     h5write.neutral(outfile, data)
 
 
-def meta(fn: Path, git_meta: dict[str, str], cfg: dict[str, T.Any]):
+def meta(fn: Path, git_meta: dict[str, str], cfg: dict[str, T.Any]) -> None:
     """
     writes JSON file with sim setup metadata
     """
@@ -134,7 +134,7 @@ def meta(fn: Path, git_meta: dict[str, str], cfg: dict[str, T.Any]):
     fn.write_text(js)
 
 
-def maggrid(filename: Path, xmag: dict[str, T.Any]):
+def maggrid(filename: Path, xmag: dict[str, T.Any]) -> None:
     filename = Path(filename).expanduser()
 
     # %% default value for gridsize
@@ -148,7 +148,4 @@ def maggrid(filename: Path, xmag: dict[str, T.Any]):
         gridsize = xmag["gridsize"]
 
     # %% write the file
-    if not filename.parent.is_dir():
-        raise FileNotFoundError(f"{filename.parent} parent directory does not exist")
-
     h5write.maggrid(filename, xmag, gridsize)