Pushing to glowpython2

Author: sunipkm

.vscode/settings.json

@@ -1,7 +1,8 @@
 {
     "autoDocstringPy.customTemplatePath": ".vscode/docstring.mustache",
     "files.associations": {
-        "*.f": "FortranFixedForm"
+        "*.f": "FortranFixedForm",
+        "*.f90": "FortranFreeForm"
     },
     "[FortranFixedForm]": {
         "editor.rulers": [

src/GLOW/rcolum.f90

@@ -22,6 +22,13 @@
 ! grazing height is lower than the bottom of the atmosphere supplied, 
 ! column densities are set to 'infinity', i.e., 1.0e30.
 
+! Input:
+!   CHI    Zenith angle (radians)
+!   ZZ     Altitude array; cm [jmax]
+!   ZMAJ   Major species number densities (O, O2, N2); cm-3 [nmaj,jmax]
+!   TN     Neutral temperature at each altitude; K [jmax]
+!   ZCOL   Slant column density (O, O2, N2); cm-2 [nmaj,jmax]
+!   ZVCD   Vertical column density (O, O2, N2) above each altitude (cumsum); cm-2 [nmaj,jmax]
 
     subroutine rcolum (chi, zz, zmaj, tn, zcol, zvcd, jmax, nmaj)
 
@@ -88,6 +95,7 @@ end subroutine rcolum
 
 !----------------------------------------------------------------------
 
+    ! Calculate Chapman Grazing Incidence Integral
     real function chap (chi, z, t, i)
 
       implicit none
@@ -97,18 +105,18 @@ real function chap (chi, z, t, i)
 
       integer,parameter ::  nmaj=3
       real,parameter :: pi=3.1415926536
-      real,parameter :: re=6.37e8
-      real,parameter :: g=978.1
+      real,parameter :: re=6.37e8 ! cm
+      real,parameter :: g=978.1 ! cm/s^2
 
       real :: am(nmaj), gr, hn, hg, hf, sqhf
       real,external :: sperfc
 
       data am/16., 32., 28./
 
-      gr=g*(re/(re+z))**2 
-      hn=1.38e-16*t/(am(i)*1.662e-24*gr)
-      hg=(re+z)/hn 
-      hf=0.5*hg*(cos(chi)**2) 
+      gr=g*(re/(re+z))**2 ! effective g
+      hn=1.38e-16*t/(am(i)*1.662e-24*gr) ! k_B * T / mass (amu * grams/amu or Dalton) * g
+      hg=(re+z)/hn ! gravitational scale height
+      hf=0.5*hg*(cos(chi)**2)
       sqhf=sqrt(hf) 
       chap=sqrt(0.5*pi*hg)*sperfc(sqhf) 
 
@@ -137,6 +145,9 @@ end function sperfc
 
 !----------------------------------------------------------------------
 
+    !Subroutine VCD
+    ! Calculates the cumulative vertical column density
+    ! for each species ZMAJ above each altitude ZZ.
     subroutine vcd(zz,zmaj,zvcd,jmax,nmaj)
 
       implicit none
@@ -155,7 +166,7 @@ subroutine vcd(zz,zmaj,zvcd,jmax,nmaj)
                        / alog(zmaj(i,jmax-1)/zmaj(i,jmax))
         do j=jmax-1,1,-1
           rat = zmaj(i,j+1) / zmaj(i,j)
-          zvcd(i,j) = zvcd(i,j+1)+zmaj(i,j)*(zz(j)-zz(j+1))/alog(rat)*(1.-rat)
+          zvcd(i,j) = zvcd(i,j+1)+zmaj(i,j)*(zz(j)-zz(j+1))/alog(rat)*(1.0-rat)
         enddo
       enddo
 

src/glowpython/rcolum.py

@@ -0,0 +1,28 @@
+# %%
+from __future__ import annotations
+from xarray import DataArray, Variable
+from typing import List, Tuple, Union, SupportsFloat as Numeric
+import numpy as np
+
+# %%
+def rcolumn(
+    major_den: DataArray,
+    t_n: DataArray,
+    chi: Numeric,
+    ) -> Tuple[DataArray, DataArray]:
+    """## Calculates the column density for each species at zenith angle `chi`.
+    First, the vertical column density is calculated, and then a fit to Chapman Grazing
+    Incidence Integral [Smith and Smith, JGR 77, 3592, 1972] is used to calculate the
+    slant column density.
+    - `chi` < 90 deg: Column densities are calculated directly.
+    - `chi` > 90 deg: Column density at grazing height for 90 deg is calculated and doubled, and the column density above the grazing height is subtracted.
+    - If grazing height is lower than the bottom of the atmosphere supplied, column densities are set to infinity (`np.inf`).
+
+    ### Args:
+        - `major_den (DataArray)`: Major species densities.
+        - `t_n (DataArray)`: Neutral temperature profile.
+        - `chi (Numeric)`: Zenith angle.
+
+    ### Returns:
+        - `Tuple[DataArray, DataArray]`: Column densities, vertical column densities.
+    """

src/glowpython/sflux.py

@@ -0,0 +1,64 @@
+# %%
+from __future__ import annotations
+from pathlib import Path
+from typing import Literal, Sequence, Tuple, Union, SupportsFloat as Numeric
+import numpy as np
+from .base import FluxSource, DATA_DIR
+from .glowfort import cglow as cg # type: ignore
+
+# %%
+class Sflux:
+    def __init__(self, kind: FluxSource | Tuple[Sequence[float], Sequence[float], Sequence[float]], data_dir: Path = DATA_DIR):
+        self._kind = kind
+        self._docalc = True
+        if kind == 'Hinteregger':
+            self._method = 'Hinteregger'
+            ifile = data_dir / 'ssflux_hint.dat'
+            self._data = np.loadtxt(ifile, skiprows=1)[::-1, :].T # reverse order
+            self._fluxmask = self._data[0] < 251 & self._data[0] > 17 # 17-251 nm
+            cg.wave1 = self._data[0]
+            cg.wave2 = self._data[1]
+        elif kind == 'EUVAC':
+            self._method = 'EUVAC'
+            ifile = data_dir / 'ssflux_euvac.dat'
+            self._data = np.loadtxt(ifile, skiprows=1)[::-1, :].T # reverse order
+            self._fluxmask = self._data[0] < 51 & self._data[0] > 1 # 1-51 nm
+            cg.wave1 = self._data[0]
+            cg.wave2 = self._data[1]
+        elif isinstance(kind, tuple):
+            self._method = 'custom'
+            self._data = np.stack(kind, axis=0)
+            if self._data.ndim != 2:
+                raise ValueError(f"Custom flux data must be 2D, got {self._data.ndim}D.")
+            if self._data.shape[1] != 3:
+                raise ValueError(f"Custom flux data must have 3 columns, got {self._data.shape[1]}.")
+            if self._data.shape[0] != 123:
+                raise ValueError(f"Custom flux data must have 123 wavelength bins, got {self._data.shape[0]}.")
+            if self._data[1, 0] > self._data[0, 0]:
+                self._data = self._data[::-1, :] # reverse order
+            cg.wave1 = self._data[:, 0]
+            cg.wave2 = self._data[:, 1]
+            self._docalc = False
+        else:
+            raise ValueError(f"Unknown flux source: {kind}. Must be 'Hinteregger', 'EUVAC', or a tuple of values.")
+        
+    def update(self, f107: Numeric, f107a: Numeric, xuvfac: int = 3) -> None:
+        if self._docalc:
+            if self._data.shape[1] == 5: # Hinteregger
+                r1 = 0.0138*(f107a - 71.5) + 0.005*(f107 - f107a + 3.9) # + 1
+                r2 = 0.59425*(f107a - 71.5) + 0.3811*(f107 - f107a + 3.9) # + 1
+                sflux = self._data[2] + r1 * self._data[3] + r2 * self._data[4]
+                np.clip(sflux, 0, None, out=sflux)
+                if xuvfac > 1e-6:
+                    sflux[self._fluxmask] *= xuvfac
+            elif self._data.shape[1] == 4: # EUVAC
+                p107 = (f107 + f107a)*0.5
+                sflux = self._data[2] * (1 + self._data[3]*(p107 - 80))
+                np.clip(sflux, 0.1*self._data[2], None, out=sflux)
+                if xuvfac > 1e-6:
+                    sflux[self._fluxmask] *= xuvfac
+            cg.sflux = sflux
+        else:
+            # Set the cg.wave1, cg.wave2, cg.sflux to these values
+            cg.sflux = self._data[:, 2]
+# %%