From b764060fa53f4c31907e6a0994fdda908f0e588c Mon Sep 17 00:00:00 2001 From: scivision Date: Tue, 4 Jun 2024 21:13:39 -0400 Subject: [PATCH] update type hints and self describe docs --- example/plot_Efield.py | 3 +- example/plotdata.py | 12 +++++- src/gemini3d/compare/plot.py | 12 ++++-- src/gemini3d/compare/utils.py | 21 +++++++++-- src/gemini3d/conductivity.py | 4 +- src/gemini3d/efield/Efield_erf.py | 10 +++-- src/gemini3d/efield/Jcurrent_gaussian.py | 18 ++++++++- src/gemini3d/efield/__init__.py | 11 +++++- src/gemini3d/hdf5/write.py | 37 ++++++++++++++++--- src/gemini3d/particles/gaussian2d.py | 19 ++++++---- src/gemini3d/plasma.py | 47 ++++++++++++++++++++---- src/gemini3d/plot/cartesian.py | 9 +++-- src/gemini3d/raw/read.py | 18 ++++++++- src/gemini3d/read.py | 32 ++++++++++++---- src/gemini3d/utils.py | 9 +++-- src/gemini3d/write.py | 18 +++++---- 16 files changed, 216 insertions(+), 64 deletions(-) diff --git a/example/plot_Efield.py b/example/plot_Efield.py index 21444b6..58f5bf0 100644 --- a/example/plot_Efield.py +++ b/example/plot_Efield.py @@ -6,13 +6,12 @@ from pathlib import Path import argparse -import xarray import matplotlib as mpl import gemini3d.read as read -def plotVmaxx1it(ax: mpl.axes.Axes, V: xarray.DataArray) -> None: +def plotVmaxx1it(ax, V) -> None: ax.set_title("Vmaxx1it: Potential") if V.ndim == 1: ax.plot(dat["mlat"], V) diff --git a/example/plotdata.py b/example/plotdata.py index 92489dc..792ee4e 100644 --- a/example/plotdata.py +++ b/example/plotdata.py @@ -1,13 +1,21 @@ #!/usr/bin/env python3 from argparse import ArgumentParser -import xarray from matplotlib.pyplot import figure, show import gemini3d.read as read -def plotplasma(dat: xarray.Dataset): +def plotplasma(dat): + """ + + Parameters + ---------- + + dat: xarray.Dataset + plasma state data to plot + + """ for k, v in dat.items(): if k == "t": t = v diff --git a/src/gemini3d/compare/plot.py b/src/gemini3d/compare/plot.py index f1e1074..15fb718 100644 --- a/src/gemini3d/compare/plot.py +++ b/src/gemini3d/compare/plot.py @@ -3,15 +3,14 @@ from datetime import datetime import typing -import xarray import matplotlib as mpl from .. import read def plotdiff( - A: xarray.DataArray, - B: xarray.DataArray, + A, + B, time: datetime, new_dir: Path, ref_dir: Path, @@ -24,6 +23,11 @@ def plotdiff( Parameters ---------- + A: xarray.DataArray + new data + B: xarray.DataArray + reference data + imax: int, optional index of maximum difference """ @@ -63,7 +67,7 @@ def plotdiff( plotdiff(A[i], B[i], time, new_dir, ref_dir, name=f"{name}-{i}") elif is3d: # pick x2 and x3 slice at maximum difference - im: dict[str, xarray.DataArray] = abs(A - B).argmax(dim=...) # type: ignore + im = abs(A - B).argmax(dim=...) ix2: int = im["x2"].data ix3: int = im["x3"].data plotdiff( diff --git a/src/gemini3d/compare/utils.py b/src/gemini3d/compare/utils.py index 365abd0..edbcd72 100644 --- a/src/gemini3d/compare/utils.py +++ b/src/gemini3d/compare/utils.py @@ -1,13 +1,26 @@ from __future__ import annotations import json -import xarray - from ..utils import get_pkg_file -def err_pct(a: xarray.DataArray, b: xarray.DataArray) -> float: - """compute maximum error percent""" +def err_pct(a, b) -> float: + """compute maximum error percent + + Parameters + ---------- + + a: xarray.DataArray + new data + b: xarray.DataArray + reference data + + Returns + ------- + + float + maximum error percent + """ return (abs(a - b).max() / abs(b).max()).item() * 100 diff --git a/src/gemini3d/conductivity.py b/src/gemini3d/conductivity.py index 22506ef..6f1edda 100644 --- a/src/gemini3d/conductivity.py +++ b/src/gemini3d/conductivity.py @@ -149,9 +149,7 @@ def collisions3D(atmos, Ts, ns, vsx1, ms): T = Ts[:, :, :, lsp - 1] nusn[:, :, :, lsp - 1, 0] = 8.9e-11 * (1 + 5.7e-4 * T) * (T**0.5) * nO * 1e-6 nusn[:, :, :, lsp - 1, 1] = 2.33e-11 * (1 - 1.21e-4 * T) * (T) * nN2 * 1e-6 - nusn[:, :, :, lsp - 1, 2] = ( - 1.82e-10 * (1 + 3.6e-2 * (T**0.5)) * (T**0.5) * nO2 * 1e-6 - ) + nusn[:, :, :, lsp - 1, 2] = 1.82e-10 * (1 + 3.6e-2 * (T**0.5)) * (T**0.5) * nO2 * 1e-6 nusn[:, :, :, lsp - 1, 3] = 4.5e-9 * (1 - 1.35e-4 * T) * (T**0.5) * nH * 1e-6 nus[:, :, :, lsp - 1] = np.sum(nusn[:, :, :, lsp - 1, :], axis=3) diff --git a/src/gemini3d/efield/Efield_erf.py b/src/gemini3d/efield/Efield_erf.py index 8d9fee3..2f1557c 100644 --- a/src/gemini3d/efield/Efield_erf.py +++ b/src/gemini3d/efield/Efield_erf.py @@ -2,22 +2,26 @@ import logging import typing as T -import xarray import numpy as np from scipy.special import erf def Efield_erf( - E: xarray.Dataset, + E, xg: dict[str, T.Any], lx1: int, lx2: int, lx3: int, gridflag: int, flagdip: bool, -) -> xarray.Dataset: +): """ synthesize a feature + + Parameters + ---------- + + E: xarray.Dataset """ if E.Etarg > 1: diff --git a/src/gemini3d/efield/Jcurrent_gaussian.py b/src/gemini3d/efield/Jcurrent_gaussian.py index 581f094..581d444 100644 --- a/src/gemini3d/efield/Jcurrent_gaussian.py +++ b/src/gemini3d/efield/Jcurrent_gaussian.py @@ -1,8 +1,22 @@ -import xarray import numpy as np -def Jcurrent_gaussian(E: xarray.Dataset, gridflag: int, flagdip: bool) -> xarray.Dataset: +def Jcurrent_gaussian(E, gridflag: int, flagdip: bool): + """ + + Parameters + ---------- + + E: xarray.Dataset + Current density + + Returns + ------- + + E: xarray.Dataset + Current density with boundary conditions + + """ if E.mlon.size > 1: shapelon = np.exp(-((E.mlon - E.mlonmean) ** 2) / 2 / E.mlonsig**2) else: diff --git a/src/gemini3d/efield/__init__.py b/src/gemini3d/efield/__init__.py index 513d16b..14e0ece 100644 --- a/src/gemini3d/efield/__init__.py +++ b/src/gemini3d/efield/__init__.py @@ -211,7 +211,16 @@ def Esigma(pwidth: float, pmax: float, pmin: float, px) -> tuple[float, T.Any]: return wsig, xsig -def check_finite(v: xarray.DataArray, name: str | None = None): +def check_finite(v, name: str | None = None): + """ + checks that the input is finite + + Parameters + ---------- + + v: xarray.DataArray + input data to check + """ i = np.logical_not(np.isfinite(v)) if i.any(): diff --git a/src/gemini3d/hdf5/write.py b/src/gemini3d/hdf5/write.py index 559473f..ed31a6a 100644 --- a/src/gemini3d/hdf5/write.py +++ b/src/gemini3d/hdf5/write.py @@ -10,14 +10,13 @@ import h5py import numpy as np -import xarray from ..utils import datetime2stem, to_datetime CLVL = 3 # GZIP compression level: larger => better compression, slower to write -def state(fn: Path, dat: xarray.Dataset) -> None: +def state(fn: Path, dat) -> None: """ write STATE VARIABLE initial conditions @@ -27,6 +26,14 @@ def state(fn: Path, dat: xarray.Dataset) -> None: INPUT ARRAYS SHOULD BE TRIMMED TO THE CORRECT SIZE I.E. THEY SHOULD NOT INCLUDE GHOST CELLS + + Parameters + ---------- + + fn: pathlib.Path + output filename + dat: xarray.Dataset + data to write """ logging.info(f"state: {fn}") @@ -42,7 +49,7 @@ def state(fn: Path, dat: xarray.Dataset) -> None: _write_var(f, "/Phiall", dat["Phitop"]) -def _write_var(fid, name: str, A: xarray.DataArray) -> None: +def _write_var(fid, name: str, A) -> None: """ NOTE: The .transpose() reverses the dimension order. The HDF Group never implemented the intended H5T_array_create(..., perm) @@ -51,6 +58,12 @@ def _write_var(fid, name: str, A: xarray.DataArray) -> None: Fortran, including the HDF Group Fortran interfaces and h5fortran as well as Matlab read/write HDF5 in Fortran order. h5py read/write HDF5 in C order so we need the .transpose() for h5py + + Parameters + ---------- + + A: xarray.DataArray + variable to write to HDF5 """ p4s = ("species", "x3", "x2", "x1") @@ -212,9 +225,15 @@ def grid(size_fn: Path, grid_fn: Path, xg: dict[str, T.Any]) -> None: h["/glatctr"] = xg["glatctr"] -def Efield(outdir: Path, E: xarray.Dataset) -> None: +def Efield(outdir: Path, E) -> None: """ write Efield to disk + + Parameters + ---------- + + E: xarray.Dataset + Electric field """ with h5py.File(outdir / "simsize.h5", "w") as f: @@ -248,7 +267,15 @@ def Efield(outdir: Path, E: xarray.Dataset) -> None: f[f"/{k}"] = E[k].loc[time].astype(np.float32) -def precip(outdir: Path, P: xarray.Dataset) -> None: +def precip(outdir: Path, P) -> None: + """ + + Parameters + ---------- + + P: xarray.Dataset + precipitation data + """ with h5py.File(outdir / "simsize.h5", "w") as f: f.create_dataset("/llon", data=P.mlon.size, dtype=np.int32) f.create_dataset("/llat", data=P.mlat.size, dtype=np.int32) diff --git a/src/gemini3d/particles/gaussian2d.py b/src/gemini3d/particles/gaussian2d.py index 3f744a7..e48b3a2 100644 --- a/src/gemini3d/particles/gaussian2d.py +++ b/src/gemini3d/particles/gaussian2d.py @@ -1,20 +1,23 @@ import numpy as np -import xarray -def gaussian2d(pg: xarray.Dataset, Qpeak: float, Qbackground: float): +def gaussian2d(pg, Qpeak: float, Qbackground: float): + """ + + Parameters + ---------- + + pg: xarray.Dataset + parameters of gaussian + """ mlon_mean = pg.mlon.mean().item() mlat_mean = pg.mlat.mean().item() if "mlon_sigma" in pg.attrs and "mlat_sigma" in pg.attrs: Q = ( Qpeak - * np.exp( - -((pg.mlon.data[:, None] - mlon_mean) ** 2) / (2 * pg.mlon_sigma**2) - ) - * np.exp( - -((pg.mlat.data[None, :] - mlat_mean) ** 2) / (2 * pg.mlat_sigma**2) - ) + * np.exp(-((pg.mlon.data[:, None] - mlon_mean) ** 2) / (2 * pg.mlon_sigma**2)) + * np.exp(-((pg.mlat.data[None, :] - mlat_mean) ** 2) / (2 * pg.mlat_sigma**2)) ) elif "mlon_sigma" in pg.attrs: Q = Qpeak * np.exp( diff --git a/src/gemini3d/plasma.py b/src/gemini3d/plasma.py index abf202e..8ac1386 100644 --- a/src/gemini3d/plasma.py +++ b/src/gemini3d/plasma.py @@ -71,9 +71,7 @@ def equilibrium_resample(p: dict[str, T.Any], xg: dict[str, T.Any]): write.state(p["indat_file"], dat_interp) -def model_resample( - xgin: dict[str, T.Any], dat: xarray.Dataset, xg: dict[str, T.Any] -) -> xarray.Dataset: +def model_resample(xgin: dict[str, T.Any], dat, xg: dict[str, T.Any]): """resample a grid usually used to upsample an equilibrium simulation grid @@ -198,7 +196,14 @@ def model_resample( return dat_interp -def check_density(n: xarray.DataArray): +def check_density(n): + """ + Parameters + ---------- + + n: xarray.DataArray + number density + """ if not np.isfinite(n).all(): raise ValueError("non-finite density") if (n < 0).any(): @@ -207,14 +212,28 @@ def check_density(n: xarray.DataArray): raise ValueError("too small maximum density") -def check_drift(v: xarray.DataArray): +def check_drift(v): + """ + Parameters + ---------- + + v: xarray.DataArray + velocity + """ if not np.isfinite(v).all(): raise ValueError("non-finite drift") if (abs(v) > 10e3).any(): raise ValueError("excessive drift velocity") -def check_temperature(Ts: xarray.DataArray): +def check_temperature(Ts): + """ + Parameters + ---------- + + Ts: xarray.DataArray + temperature + """ if not np.isfinite(Ts).all(): raise ValueError("non-finite temperature") if (Ts < 0).any(): @@ -223,12 +242,26 @@ def check_temperature(Ts: xarray.DataArray): raise ValueError("too cold maximum temperature") -def equilibrium_state(p: dict[str, T.Any], xg: dict[str, T.Any]) -> xarray.Dataset: +def equilibrium_state(p: dict[str, T.Any], xg: dict[str, T.Any]): """ generate (arbitrary) initial conditions for a grid. NOTE: only works on symmmetric closed grids! [f107a, f107, ap] = activ + + Parameters + ---------- + + p: dict + simulation parameters + xg: dict + simulation grid + + Returns + ------- + + dat: xarray.Dataset + initial conditions of equilibrium state """ # %% MAKE UP SOME INITIAL CONDITIONS FOR FORTRAN CODE diff --git a/src/gemini3d/plot/cartesian.py b/src/gemini3d/plot/cartesian.py index 0bf6f28..0420a0b 100644 --- a/src/gemini3d/plot/cartesian.py +++ b/src/gemini3d/plot/cartesian.py @@ -3,9 +3,9 @@ from pathlib import Path import typing as T from datetime import datetime -import numpy as np import math -import xarray + +import numpy as np import scipy.interpolate as interp from matplotlib.figure import Figure @@ -29,7 +29,7 @@ def plot_interp( fg: Figure, time: datetime, xg: dict[str, T.Any], - parm: xarray.DataArray, + parm, *, name: str = "", ref_alt: float = REF_ALT, @@ -40,6 +40,9 @@ def plot_interp( Parameters ---------- + parm: xarray.DataArray + parameter to plot + xp: eastward distance (rads.) should be interpreted as northward distance (in rads.). Irrespective of ordering of xg.theta, this will be monotonic increasing!!! diff --git a/src/gemini3d/raw/read.py b/src/gemini3d/raw/read.py index f890591..f74a863 100644 --- a/src/gemini3d/raw/read.py +++ b/src/gemini3d/raw/read.py @@ -331,7 +331,15 @@ def frame3d_curvavg(file: Path, xg: dict[str, T.Any] | None = None) -> xarray.Da return dat -def frame3d_curvne(file: Path, xg: dict[str, T.Any] | None = None) -> xarray.Dataset: +def frame3d_curvne(file: Path, xg: dict[str, T.Any] | None = None): + """ + + Returns + ------- + + dat: xarray.Dataset + 3D simulation data + """ if not file.is_file(): raise FileNotFoundError(file) @@ -394,9 +402,15 @@ def read2D(f: T.BinaryIO, lx: tuple[int, ...] | list[int]): raise ValueError(f"lx must have 2 or 3 elements, you have lx={lx}") -def glow_aurmap(file: Path, xg: dict[str, T.Any] | None = None) -> xarray.Dataset: +def glow_aurmap(file: Path, xg: dict[str, T.Any] | None = None): """ read the auroral output from GLOW + + Returns + ------- + + dat: xarray.Dataset + GLOW auroral data """ lx = simsize(file.parent) diff --git a/src/gemini3d/read.py b/src/gemini3d/read.py index d1eadcb..44f5a44 100644 --- a/src/gemini3d/read.py +++ b/src/gemini3d/read.py @@ -133,9 +133,21 @@ def frame( return dat -def derive(dat: xarray.Dataset, var: set[str], flag: int) -> xarray.Dataset: +def derive(dat, var: set[str], flag: int): """ compute derived variables based on file data + + Parameters + ---------- + + dat: xarray.DataSet + data to derive from + + Returns + ------- + + dat: xarray.DataSet + derived data """ lx = (dat.sizes["x1"], dat.sizes["x2"], dat.sizes["x3"]) @@ -174,12 +186,18 @@ def derive(dat: xarray.Dataset, var: set[str], flag: int) -> xarray.Dataset: return dat -def glow(fn: Path) -> xarray.Dataset: - """read GLOW data""" +def glow(fn: Path): + """read GLOW data + + Returns + ------- + + xarray.Dataset + """ return h5read.glow_aurmap(fn) -def Efield(fn: Path) -> xarray.Dataset: +def Efield(fn: Path): """load Efield data "Efield_inputs" Parameters @@ -189,7 +207,7 @@ def Efield(fn: Path) -> xarray.Dataset: Returns ------- - dat: dict of np.ndarray + dat: xarray.Dataset electric field """ @@ -198,7 +216,7 @@ def Efield(fn: Path) -> xarray.Dataset: return h5read.Efield(fn) -def precip(fn: Path) -> xarray.Dataset: +def precip(fn: Path): """load precipitation to disk Parameters @@ -208,7 +226,7 @@ def precip(fn: Path) -> xarray.Dataset: Returns ------- - dat: dict + dat: xarray.Dataset precipitation """ diff --git a/src/gemini3d/utils.py b/src/gemini3d/utils.py index e85de6a..cdd0d75 100644 --- a/src/gemini3d/utils.py +++ b/src/gemini3d/utils.py @@ -48,7 +48,7 @@ def str2func(name: str, path: Path | None = None) -> T.Callable: 0. file in "path" (if present) 1. os.getcwd()/name.py containing function name() - 2. gemiin3d. /name.py module file containing function name() + 2. gemini3d. /name.py module file containing function name() 3. gemini3d. /__init__.py containing function name() @@ -90,16 +90,17 @@ def str2func(name: str, path: Path | None = None) -> T.Callable: return getattr(mod, func_name) -def to_datetime(times: xarray.DataArray | np.datetime64 | datetime) -> datetime: +def to_datetime(times) -> datetime: """ Parameters ---------- - atimes : xarray time + atimes: xarray.DataArray, np.datetime64, datetime.datetime + time Returns ------- - times : list[datetime.datetime] + times: datetime.datetime """ if isinstance(times, datetime): diff --git a/src/gemini3d/write.py b/src/gemini3d/write.py index ad08bf9..c48a326 100644 --- a/src/gemini3d/write.py +++ b/src/gemini3d/write.py @@ -5,13 +5,11 @@ import logging import json -import xarray - from .utils import git_meta from .hdf5 import write as h5write -def state(out_file: Path, dat: xarray.Dataset, **kwargs) -> None: +def state(out_file: Path, dat, **kwargs) -> None: """ WRITE STATE VARIABLE DATA. NOTE: WE don't write ANY OF THE ELECTRODYNAMIC @@ -20,6 +18,12 @@ def state(out_file: Path, dat: xarray.Dataset, **kwargs) -> None: INPUT ARRAYS SHOULD BE TRIMMED TO THE CORRECT SIZE I.E. THEY SHOULD NOT INCLUDE GHOST CELLS + + Parameters + ---------- + + dat: xarray.Dataset + state variables to write """ # %% allow overriding "dat" @@ -59,13 +63,13 @@ def grid(cfg: dict[str, T.Any], xg: dict[str, T.Any]) -> None: meta(input_dir / "setup_grid.json", git_meta(), cfg) -def Efield(E: xarray.Dataset, outdir: Path) -> None: +def Efield(E, outdir: Path) -> None: """writes E-field to disk Parameters ---------- - E: dict + E: xarray.Dataset E-field values outdir: pathlib.Path directory to write files into @@ -77,12 +81,12 @@ def Efield(E: xarray.Dataset, outdir: Path) -> None: h5write.Efield(outdir, E) -def precip(precip: xarray.Dataset, outdir: Path) -> None: +def precip(precip, outdir: Path) -> None: """writes precipitation to disk Parameters ---------- - precip: dict + precip: xarray.Dataset preicipitation values outdir: pathlib.Path directory to write files into