diff --git a/.gitignore b/.gitignore index ebebca9..3aca6e2 100644 --- a/.gitignore +++ b/.gitignore @@ -75,3 +75,7 @@ output # cache *__pycache__* + +# release/debug +*.release +*.debug diff --git a/README.md b/README.md index 776fcec..91e0765 100644 --- a/README.md +++ b/README.md @@ -110,9 +110,9 @@ written on `(del_temperature, pressure, wavenumber)`, include a `pressure` coordinate in Pa, and still retain degenerate `del_temperature` and `pressure` dimensions when only one state is requested. -Across `pyharp.spectra`, wavenumber ranges are interpreted as -lower-inclusive and upper-exclusive: `--wn-range=20,22` with `1 cm^-1` -resolution samples `20` and `21`, not `22`. +Across `pyharp.spectra`, wavenumber ranges are interpreted as inclusive at +both ends: `--wn-range=20,22` with `1 cm^-1` resolution samples `20`, `21`, +and `22`. See the [pyharp-dump CLI documentation](https://pyharp.readthedocs.io/en/latest/dump_cli.html) for the full command reference, output naming conventions, and NetCDF schema. @@ -133,10 +133,11 @@ pyharp-plot overview --species H2O CO2 --temperature-k 300 --pressure-bar 1 --wn Use `--pair` for CIA pairs, `--species` for molecules, and `--composition` for gas mixtures such as `H2O:0.1,H2:0.9`. All plot commands accept `--wn-range=min,max`; `overview` accepts multiple `--wn-range` values for -multi-page PDFs. These ranges are lower-inclusive and upper-exclusive. Use -`--output` to choose the output path. Without `--output`, plots are written -under `--output-dir` (default `output/`) with names derived from the target, -plot type, temperature, pressure, and wavenumber range. +multi-page PDFs. These ranges include both endpoints, so adjacent repeated +ranges such as `20,2500` and `2500,10000` both include `2500`. Use `--output` +to choose the output path. Without `--output`, plots are written under +`--output-dir` (default `output/`) with names derived from the target, plot +type, temperature, pressure, and wavenumber range. For plot commands that use pressure, `--temperature-k` and `--pressure-bar` also accept matched comma-separated vectors such as `--temperature-k 300,400 --pressure-bar 1,10`. `pyharp-plot` then runs one diff --git a/docs/source/builtin_opacities.rst b/docs/source/builtin_opacities.rst index ada0d70..9793b7e 100644 --- a/docs/source/builtin_opacities.rst +++ b/docs/source/builtin_opacities.rst @@ -51,6 +51,16 @@ The following functions have been used to process the legacy CIA data files: :undoc-members: :imported-members: +NetCDF dump-backed molecular opacities +-------------------------------------- + +Pyharp also supports NetCDF dump-backed molecular absorption through the +``molecule-line`` and ``cia`` opacity types. These readers consume +cross-section fields on ``(del_temperature, pressure, wavenumber)`` or any +equivalent variable-dimension ordering, convert them into the runtime units +used by the C++ radiative-transfer core, and interpolate on the requested +band grid. + References ---------- .. [1] Lupu, R., et al. "Correlated k coefficients for H2-He atmospheres; 196 spectral windows and 1460 pressure-temperature points." Zenodo, doi 0.5281/zenodo.5590988 (2021). diff --git a/docs/source/dump_cli.rst b/docs/source/dump_cli.rst index 88d96a6..62c5487 100644 --- a/docs/source/dump_cli.rst +++ b/docs/source/dump_cli.rst @@ -75,8 +75,8 @@ Shared Options Wavenumber bounds in ``cm^-1``. Repeat this option to write one NetCDF file per band. When ``--output`` is provided with multiple ranges, pyharp appends ``__`` to the requested stem. Ranges are - lower-inclusive and upper-exclusive, so ``--wn-range=20,22`` with - ``--resolution 1`` samples ``20`` and ``21``. + inclusive at both ends, so ``--wn-range=20,22`` with ``--resolution 1`` + samples ``20``, ``21``, and ``22``. ``--resolution value`` Wavenumber spacing in ``cm^-1``. The default is ``1``. @@ -177,8 +177,8 @@ Repeat ``--wn-range`` to compute multiple bands in one run: pyharp-dump xsection --species H2O --temperature-k 300 --pressure-bar 1 \ --wn-range=20,2500 --wn-range=2500,10000 -This writes one file for ``[20, 2500)`` and one file for ``[2500, 10000)``. -Adjacent repeated ranges do not duplicate the boundary sample. If +This writes one file for ``[20, 2500]`` and one file for ``[2500, 10000]``. +Adjacent repeated ranges duplicate the shared boundary sample. If ``--output output/h2o.nc`` is provided, the generated files are ``output/h2o_20_2500.nc`` and ``output/h2o_2500_10000.nc``. diff --git a/docs/source/opacity.rst b/docs/source/opacity.rst index 7179624..b42c695 100644 --- a/docs/source/opacity.rst +++ b/docs/source/opacity.rst @@ -23,12 +23,12 @@ A complete list of built-in opacity types is given in the table below. * - 'jit' - '.pt' (saved by :func:`torch.jit.save` or :func:`pyharp.compile`) - Just-In-Time scripted opacity model - * - 'rfm-lbl' + * - 'molecule-line' - NetCDF - - Line-by-line absorption data computed by RFM - * - 'rfm-ck' + - Molecular line cross sections on ``(wavenumber, pressure, del_temperature)`` with optional same-species continuum fields + * - 'molecule-cia' - NetCDF - - Correlated-k absorption computed from line-by-line data + - Collision-induced absorption binary coefficients on ``(wavenumber, pressure, del_temperature)`` * - 'multiband-ck' - '.pt' (saved by :func:`torch.jit.save`) - Three-dimensional correlated-k opacity lookup table. The axis are @@ -52,6 +52,13 @@ A complete list of built-in opacity types is given in the table below. We give a few examples showing how to use these opacity class to load and compute optical properties. +The NetCDF dump-backed molecular opacities expect the coordinate variables +``wavenumber``, ``pressure``, ``del_temperature``, and ``temperature``. +Line opacity variables follow the naming convention ``sigma_line_``. +Optional same-species continuum terms use ``sigma_continuum__*`` and +are summed into the line attenuation automatically. CIA variables use +``binary_absorption_coefficient__``. + .. _example_sonora: Example 1. Compute Sonora2020 molecular opacities diff --git a/docs/source/opacity_classes.rst b/docs/source/opacity_classes.rst index 0d8eb4e..f69d5c6 100644 --- a/docs/source/opacity_classes.rst +++ b/docs/source/opacity_classes.rst @@ -15,7 +15,10 @@ Opacity Classes .. autoclass:: pyharp.opacity.WaveTemp :members: __init__, forward -.. autoclass:: pyharp.opacity.RFM +.. autoclass:: pyharp.opacity.MoleculeLine + :members: __init__, forward + +.. autoclass:: pyharp.opacity.MoleculeCIA :members: __init__, forward .. autoclass:: pyharp.opacity.JITOpacity diff --git a/docs/source/plot_cli.rst b/docs/source/plot_cli.rst index 641bf84..38eb38a 100644 --- a/docs/source/plot_cli.rst +++ b/docs/source/plot_cli.rst @@ -116,8 +116,8 @@ Shared Options Wavenumber bounds in ``cm^-1``. CIA plots default to ``20,10000`` when it is omitted. Molecular and mixture plots default to ``20,2500``. The ``overview`` subcommand accepts this option more than once. Ranges are - lower-inclusive and upper-exclusive, so ``--wn-range=20,22`` with - ``--resolution 1`` samples ``20`` and ``21``. + inclusive at both ends, so ``--wn-range=20,22`` with ``--resolution 1`` + samples ``20``, ``21``, and ``22``. ``--resolution value`` Wavenumber grid spacing in ``cm^-1``. The default is ``1``. diff --git a/examples/2025-amars/CMakeLists.txt b/examples/2025-amars/CMakeLists.txt deleted file mode 100644 index 8202c19..0000000 --- a/examples/2025-amars/CMakeLists.txt +++ /dev/null @@ -1,12 +0,0 @@ - -#setup_problem(amars_sw) -setup_problem(amars_lw) -#setup_problem(amars_rt_cli) -#setup_problem(amars_rt) -#setup_problem(amars_rt_rk3) - -# 4. Copy input files to run directory -file(GLOB inputs *.inp *.yaml *.txt *.atm) -foreach(input ${inputs}) - execute_process(COMMAND ln -sf ${input} ${CMAKE_BINARY_DIR}/bin/${inp}) -endforeach() diff --git a/examples/2025-amars/TVals.txt b/examples/2025-amars/TVals.txt deleted file mode 100644 index c22679a..0000000 --- a/examples/2025-amars/TVals.txt +++ /dev/null @@ -1,40 +0,0 @@ -2.100000000000000000e+02 -2.095408374643015463e+02 -2.090337589662874223e+02 -2.084737642298763944e+02 -2.078553311745680503e+02 -2.071723614624789889e+02 -2.064181203629335641e+02 -2.055851703416164753e+02 -2.046652977194134735e+02 -2.036494316777263123e+02 -2.025275548115775734e+02 -2.012886043484737684e+02 -1.999203630589509828e+02 -1.984093387830771462e+02 -1.967406313849278376e+02 -1.948977858230801132e+02 -1.928626298882599883e+02 -1.906150950080821929e+02 -1.881330183518468573e+02 -1.853919242839581329e+02 -1.823647830108876633e+02 -1.790217440417125374e+02 -1.753298418338953297e+02 -1.712526707216935904e+02 -1.667500259216838856e+02 -1.617775070753737623e+02 -1.562860804194545210e+02 -1.562860804194545210e+02 -1.562860804194545210e+02 -1.562860804194545210e+02 -1.562860804194545210e+02 -1.562860804194545210e+02 -1.585501251820229811e+02 -1.610504346114359180e+02 -1.638116641357670176e+02 -1.668610421032884119e+02 -1.702286382798712054e+02 -1.806419001195335738e+02 -1.921418391662849103e+02 -2.048418557523733909e+02 diff --git a/examples/2025-amars/aerosol_output_data.txt b/examples/2025-amars/aerosol_output_data.txt deleted file mode 100644 index b9865ed..0000000 --- a/examples/2025-amars/aerosol_output_data.txt +++ /dev/null @@ -1,201 +0,0 @@ -# p [bar] T[K] SO4AER[mol/mol] S8AER[mol/mol] -4.873999999999999999e-01 2.000000000000000000e+02 2.291999999999999877e-07 7.581000000000000187e-09 -4.637000000000000011e-01 2.000000000000000000e+02 2.462999999999999754e-07 8.089999999999999981e-09 -4.410999999999999921e-01 2.000000000000000000e+02 2.772999999999999751e-07 9.050999999999999454e-09 -4.197000000000000175e-01 2.000000000000000000e+02 3.242000000000000046e-07 1.050999999999999948e-08 -3.992999999999999883e-01 2.000000000000000000e+02 3.858999999999999814e-07 1.240999999999999995e-08 -3.798000000000000265e-01 2.000000000000000000e+02 4.597000000000000202e-07 1.466000000000000025e-08 -3.613999999999999990e-01 2.000000000000000000e+02 5.444000000000000308e-07 1.719999999999999851e-08 -3.437999999999999945e-01 2.000000000000000000e+02 6.392999999999999608e-07 2.000000000000000042e-08 -3.271000000000000019e-01 2.000000000000000000e+02 7.438999999999999842e-07 2.300999999999999892e-08 -3.111999999999999766e-01 2.000000000000000000e+02 8.578999999999999727e-07 2.619999999999999969e-08 -2.960999999999999743e-01 2.000000000000000000e+02 9.807999999999999946e-07 2.953000000000000039e-08 -2.817000000000000060e-01 2.000000000000000000e+02 1.111999999999999899e-06 3.295999999999999868e-08 -2.681000000000000050e-01 2.000000000000000000e+02 1.252000000000000041e-06 3.641999999999999790e-08 -2.550999999999999934e-01 2.000000000000000000e+02 1.397999999999999998e-06 3.986000000000000091e-08 -2.426999999999999991e-01 2.000000000000000000e+02 1.551000000000000058e-06 4.332999999999999823e-08 -2.308999999999999941e-01 2.000000000000000000e+02 1.709000000000000070e-06 4.701000000000000207e-08 -2.197000000000000064e-01 2.000000000000000000e+02 1.870000000000000095e-06 5.085000000000000204e-08 -2.091000000000000081e-01 2.000000000000000000e+02 2.031000000000000121e-06 5.477000000000000008e-08 -1.990000000000000102e-01 2.000000000000000000e+02 2.190000000000000208e-06 5.871000000000000095e-08 -1.892999999999999960e-01 2.000000000000000000e+02 2.352000000000000097e-06 6.274000000000000461e-08 -1.801999999999999991e-01 2.000000000000000000e+02 2.522999999999999815e-06 6.700000000000000438e-08 -1.713999999999999968e-01 2.000000000000000000e+02 2.726999999999999895e-06 7.204000000000000186e-08 -1.630999999999999950e-01 2.000000000000000000e+02 3.008999999999999905e-06 7.886000000000000602e-08 -1.552000000000000046e-01 2.000000000000000000e+02 3.330999999999999956e-06 8.622000000000000046e-08 -1.476999999999999980e-01 2.000000000000000000e+02 3.532999999999999886e-06 8.978000000000000057e-08 -1.406000000000000028e-01 2.000000000000000000e+02 3.413999999999999839e-06 8.570000000000000639e-08 -1.338000000000000023e-01 2.000000000000000000e+02 2.928999999999999824e-06 7.978000000000000367e-08 -1.272999999999999965e-01 2.000000000000000000e+02 2.347999999999999796e-06 7.249999999999999407e-08 -1.212000000000000022e-01 2.000000000000000000e+02 1.841000000000000034e-06 6.499999999999999971e-08 -1.152999999999999997e-01 2.000000000000000000e+02 1.435000000000000025e-06 5.779000000000000330e-08 -1.097999999999999948e-01 2.000000000000000000e+02 1.111000000000000035e-06 5.086000000000000014e-08 -1.044999999999999957e-01 2.000000000000000000e+02 8.591999999999999645e-07 4.443000000000000146e-08 -9.941999999999999449e-02 2.000000000000000000e+02 6.670000000000000302e-07 3.864000000000000057e-08 -9.463000000000000578e-02 2.000000000000000000e+02 5.232000000000000263e-07 3.358999999999999836e-08 -9.006000000000000116e-02 2.000000000000000000e+02 4.162999999999999889e-07 2.928999999999999956e-08 -8.572000000000000453e-02 2.000000000000000000e+02 3.357999999999999761e-07 2.563999999999999997e-08 -8.158999999999999586e-02 2.000000000000000000e+02 2.740999999999999993e-07 2.253999999999999868e-08 -7.766000000000000680e-02 2.000000000000000000e+02 2.264999999999999966e-07 1.989999999999999952e-08 -7.391999999999999960e-02 2.000000000000000000e+02 1.896000000000000037e-07 1.769000000000000157e-08 -7.036000000000000587e-02 2.000000000000000000e+02 1.608999999999999916e-07 1.584999999999999965e-08 -6.697000000000000175e-02 2.000000000000000000e+02 1.383999999999999952e-07 1.431000000000000042e-08 -6.375000000000000111e-02 2.000000000000000000e+02 1.205000000000000004e-07 1.302000000000000012e-08 -6.067999999999999783e-02 2.000000000000000000e+02 1.059000000000000020e-07 1.192000000000000020e-08 -5.775999999999999884e-02 2.000000000000000000e+02 9.390000000000000041e-08 1.097999999999999972e-08 -5.498999999999999722e-02 2.000000000000000000e+02 8.398999999999999968e-08 1.016999999999999942e-08 -5.233999999999999764e-02 2.000000000000000000e+02 7.569999999999999626e-08 9.479000000000000044e-09 -4.982999999999999929e-02 2.000000000000000000e+02 6.869999999999999975e-08 8.877000000000000675e-09 -4.742999999999999994e-02 2.000000000000000000e+02 6.271999999999999517e-08 8.348999999999999190e-09 -4.515000000000000263e-02 2.000000000000000000e+02 5.756999999999999869e-08 7.885000000000000792e-09 -4.299000000000000044e-02 2.000000000000000000e+02 5.309999999999999903e-08 7.473000000000000147e-09 -4.091999999999999804e-02 2.000000000000000000e+02 4.918999999999999909e-08 7.106000000000000070e-09 -3.896000000000000157e-02 2.000000000000000000e+02 4.572999999999999987e-08 6.775999999999999762e-09 -3.708999999999999797e-02 2.000000000000000000e+02 4.265000000000000141e-08 6.476000000000000384e-09 -3.531000000000000111e-02 2.000000000000000000e+02 3.989999999999999994e-08 6.201999999999999883e-09 -3.361999999999999711e-02 2.000000000000000000e+02 3.742000000000000023e-08 5.950000000000000340e-09 -3.200999999999999679e-02 2.000000000000000000e+02 3.518999999999999946e-08 5.720000000000000100e-09 -3.047000000000000056e-02 2.000000000000000000e+02 3.319000000000000140e-08 5.508999999999999898e-09 -2.901000000000000106e-02 2.000000000000000000e+02 3.134999999999999948e-08 5.311000000000000375e-09 -2.761999999999999872e-02 2.000000000000000000e+02 2.965999999999999891e-08 5.124000000000000090e-09 -2.630000000000000046e-02 2.000000000000000000e+02 2.808999999999999874e-08 4.946000000000000084e-09 -2.503999999999999976e-02 2.000000000000000000e+02 2.665000000000000041e-08 4.779000000000000144e-09 -2.384000000000000008e-02 2.000000000000000000e+02 2.532000000000000107e-08 4.623000000000000269e-09 -2.270000000000000143e-02 2.000000000000000000e+02 2.410000000000000074e-08 4.477000000000000153e-09 -2.162000000000000033e-02 2.000000000000000000e+02 2.298000000000000129e-08 4.339000000000000009e-09 -2.058000000000000107e-02 2.000000000000000000e+02 2.193999999999999992e-08 4.205999999999999745e-09 -1.959999999999999937e-02 2.000000000000000000e+02 2.096000000000000041e-08 4.078999999999999666e-09 -1.865999999999999950e-02 2.000000000000000000e+02 2.005999999999999897e-08 3.955999999999999987e-09 -1.777000000000000107e-02 2.000000000000000000e+02 1.921999999999999939e-08 3.838000000000000188e-09 -1.692000000000000101e-02 2.000000000000000000e+02 1.843999999999999836e-08 3.724999999999999855e-09 -1.610999999999999932e-02 2.000000000000000000e+02 1.771000000000000109e-08 3.616999999999999814e-09 -1.535000000000000080e-02 2.000000000000000000e+02 1.703000000000000096e-08 3.511999999999999866e-09 -1.460999999999999972e-02 2.000000000000000000e+02 1.640000000000000127e-08 3.411999999999999798e-09 -1.392000000000000008e-02 2.000000000000000000e+02 1.581999999999999872e-08 3.317999999999999916e-09 -1.324999999999999963e-02 2.000000000000000000e+02 1.528999999999999993e-08 3.228999999999999913e-09 -1.261999999999999927e-02 2.000000000000000000e+02 1.479000000000000042e-08 3.143999999999999896e-09 -1.201999999999999943e-02 2.000000000000000000e+02 1.432999999999999994e-08 3.061999999999999973e-09 -1.145000000000000011e-02 2.000000000000000000e+02 1.389000000000000063e-08 2.984000000000000035e-09 -1.089999999999999997e-02 2.000000000000000000e+02 1.349000000000000036e-08 2.910000000000000084e-09 -1.038000000000000034e-02 2.000000000000000000e+02 1.310999999999999960e-08 2.840000000000000119e-09 -9.887999999999999221e-03 2.000000000000000000e+02 1.275999999999999978e-08 2.774000000000000140e-09 -9.417000000000000023e-03 2.000000000000000000e+02 1.242999999999999947e-08 2.710999999999999841e-09 -8.968999999999999542e-03 2.000000000000000000e+02 1.212000000000000033e-08 2.651000000000000048e-09 -8.543000000000000246e-03 2.000000000000000000e+02 1.181999999999999930e-08 2.593999999999999934e-09 -8.135999999999999135e-03 2.000000000000000000e+02 1.154999999999999920e-08 2.539999999999999914e-09 -7.749999999999999944e-03 2.000000000000000000e+02 1.129000000000000051e-08 2.488999999999999987e-09 -7.381000000000000408e-03 2.000000000000000000e+02 1.104999999999999969e-08 2.441000000000000153e-09 -7.030999999999999923e-03 2.000000000000000000e+02 1.082000000000000027e-08 2.395999999999999998e-09 -6.696999999999999828e-03 2.000000000000000000e+02 1.061000000000000038e-08 2.353999999999999936e-09 -6.379000000000000122e-03 2.000000000000000000e+02 1.042000000000000000e-08 2.314000000000000074e-09 -6.075999999999999804e-03 2.000000000000000000e+02 1.022999999999999962e-08 2.276999999999999892e-09 -5.786999999999999610e-03 2.000000000000000000e+02 1.005000000000000066e-08 2.241000000000000016e-09 -5.512999999999999672e-03 2.000000000000000000e+02 9.888000000000000596e-09 2.207999999999999820e-09 -5.250999999999999591e-03 2.000000000000000000e+02 9.733999999999999680e-09 2.175999999999999930e-09 -5.002000000000000370e-03 2.000000000000000000e+02 9.590000000000000177e-09 2.147000000000000134e-09 -4.765000000000000138e-03 2.000000000000000000e+02 9.454999999999999300e-09 2.120000000000000124e-09 -4.538999999999999632e-03 2.000000000000000000e+02 9.326000000000000262e-09 2.093000000000000114e-09 -4.323999999999999719e-03 2.000000000000000000e+02 9.203999999999999236e-09 2.067999999999999890e-09 -4.120000000000000398e-03 2.000000000000000000e+02 9.088000000000000050e-09 2.043000000000000079e-09 -3.924000000000000404e-03 2.000000000000000000e+02 8.979000000000000530e-09 2.019000000000000162e-09 -3.739000000000000136e-03 2.000000000000000000e+02 8.875999999999999541e-09 1.994999999999999832e-09 -3.562000000000000062e-03 2.000000000000000000e+02 8.777999999999999259e-09 1.970999999999999915e-09 -3.393000000000000183e-03 2.000000000000000000e+02 8.684999999999999684e-09 1.947999999999999891e-09 -3.233000000000000197e-03 2.000000000000000000e+02 8.595000000000000201e-09 1.923999999999999973e-09 -3.079999999999999839e-03 2.000000000000000000e+02 8.506999999999999678e-09 1.900999999999999949e-09 -2.933999999999999976e-03 2.000000000000000000e+02 8.419000000000000809e-09 1.877999999999999925e-09 -2.795999999999999874e-03 2.000000000000000000e+02 8.331999999999999765e-09 1.854999999999999901e-09 -2.664000000000000135e-03 2.000000000000000000e+02 8.245999999999999856e-09 1.832000000000000084e-09 -2.537999999999999892e-03 2.000000000000000000e+02 8.159999999999999946e-09 1.809999999999999953e-09 -2.418000000000000010e-03 2.000000000000000000e+02 8.074999999999999515e-09 1.788000000000000029e-09 -2.304000000000000058e-03 2.000000000000000000e+02 7.991999999999999699e-09 1.765999999999999899e-09 -2.194999999999999903e-03 2.000000000000000000e+02 7.908999999999999882e-09 1.743999999999999975e-09 -2.092000000000000109e-03 2.000000000000000000e+02 7.828000000000000679e-09 1.722999999999999944e-09 -1.992999999999999980e-03 2.000000000000000000e+02 7.749000000000000434e-09 1.703000000000000013e-09 -1.899000000000000080e-03 2.000000000000000000e+02 7.672000000000000804e-09 1.681999999999999982e-09 -1.809000000000000061e-03 2.000000000000000000e+02 7.596000000000000653e-09 1.662999999999999944e-09 -1.724000000000000055e-03 2.000000000000000000e+02 7.521999999999999461e-09 1.643999999999999906e-09 -1.642999999999999929e-03 2.000000000000000000e+02 7.451000000000000016e-09 1.625000000000000075e-09 -1.565999999999999901e-03 2.000000000000000000e+02 7.382000000000000358e-09 1.608000000000000031e-09 -1.492000000000000054e-03 2.000000000000000000e+02 7.314000000000000179e-09 1.590000000000000093e-09 -1.422000000000000087e-03 2.000000000000000000e+02 7.249000000000000093e-09 1.573000000000000048e-09 -1.355000000000000085e-03 2.000000000000000000e+02 7.185999999999999794e-09 1.557000000000000103e-09 -1.291000000000000047e-03 2.000000000000000000e+02 7.125000000000000108e-09 1.542000000000000052e-09 -1.229999999999999974e-03 2.000000000000000000e+02 7.066000000000000208e-09 1.527000000000000000e-09 -1.172000000000000082e-03 2.000000000000000000e+02 7.009000000000000095e-09 1.511999999999999949e-09 -1.116999999999999937e-03 2.000000000000000000e+02 6.953999999999999768e-09 1.497999999999999997e-09 -1.064999999999999974e-03 2.000000000000000000e+02 6.901000000000000054e-09 1.484999999999999938e-09 -1.015000000000000060e-03 2.000000000000000000e+02 6.848999999999999820e-09 1.470999999999999987e-09 -9.671000000000000343e-04 2.000000000000000000e+02 6.799000000000000199e-09 1.459000000000000028e-09 -9.217000000000000149e-04 2.000000000000000000e+02 6.751000000000000365e-09 1.445999999999999970e-09 -8.785000000000000491e-04 2.000000000000000000e+02 6.704000000000000011e-09 1.434999999999999904e-09 -8.371999999999999668e-04 2.000000000000000000e+02 6.659000000000000269e-09 1.422999999999999946e-09 -7.979999999999999897e-04 2.000000000000000000e+02 6.613999999999999701e-09 1.412000000000000087e-09 -7.604999999999999998e-04 2.000000000000000000e+02 6.570999999999999746e-09 1.401000000000000022e-09 -7.249000000000000018e-04 2.000000000000000000e+02 6.530000000000000405e-09 1.391000000000000056e-09 -6.909000000000000427e-04 2.000000000000000000e+02 6.489000000000000236e-09 1.379999999999999991e-09 -6.585000000000000141e-04 2.000000000000000000e+02 6.449000000000000374e-09 1.370000000000000025e-09 -6.277000000000000244e-04 2.000000000000000000e+02 6.408999999999999685e-09 1.360000000000000060e-09 -5.982999999999999603e-04 2.000000000000000000e+02 6.368000000000000344e-09 1.350000000000000094e-09 -5.703000000000000387e-04 2.000000000000000000e+02 6.324000000000000082e-09 1.339000000000000029e-09 -5.435999999999999860e-04 2.000000000000000000e+02 6.279000000000000341e-09 1.327999999999999964e-09 -5.182000000000000191e-04 2.000000000000000000e+02 6.231999999999999986e-09 1.316000000000000005e-09 -4.939000000000000248e-04 2.000000000000000000e+02 6.182999999999999845e-09 1.304000000000000046e-09 -4.708000000000000054e-04 2.000000000000000000e+02 6.131999999999999918e-09 1.290999999999999988e-09 -4.488000000000000127e-04 2.000000000000000000e+02 6.079000000000000204e-09 1.277999999999999929e-09 -4.277999999999999902e-04 2.000000000000000000e+02 6.023999999999999877e-09 1.263999999999999978e-09 -4.077999999999999919e-04 2.000000000000000000e+02 5.966999999999999764e-09 1.250000000000000026e-09 -3.888000000000000180e-04 2.000000000000000000e+02 5.907000000000000385e-09 1.234999999999999975e-09 -3.706000000000000093e-04 2.000000000000000000e+02 5.845999999999999872e-09 1.219999999999999923e-09 -3.533000000000000224e-04 2.000000000000000000e+02 5.782000000000000093e-09 1.203999999999999978e-09 -3.368000000000000009e-04 2.000000000000000000e+02 5.715999999999999700e-09 1.188000000000000033e-09 -3.210999999999999987e-04 2.000000000000000000e+02 5.647000000000000042e-09 1.170999999999999989e-09 -3.061000000000000136e-04 2.000000000000000000e+02 5.575999999999999770e-09 1.153999999999999944e-09 -2.918999999999999937e-04 2.000000000000000000e+02 5.502000000000000232e-09 1.136000000000000006e-09 -2.782999999999999884e-04 2.000000000000000000e+02 5.426000000000000081e-09 1.116999999999999968e-09 -2.652999999999999976e-04 2.000000000000000000e+02 5.346999999999999837e-09 1.097999999999999931e-09 -2.529000000000000215e-04 2.000000000000000000e+02 5.265999999999999806e-09 1.078000000000000000e-09 -2.412000000000000081e-04 2.000000000000000000e+02 5.181999999999999683e-09 1.058000000000000069e-09 -2.299000000000000045e-04 2.000000000000000000e+02 5.095000000000000294e-09 1.037000000000000038e-09 -2.191999999999999884e-04 2.000000000000000000e+02 5.006000000000000291e-09 1.016000000000000007e-09 -2.090000000000000115e-04 2.000000000000000000e+02 4.914000000000000195e-09 9.936999999999999910e-10 -1.992999999999999926e-04 2.000000000000000000e+02 4.820000000000000313e-09 9.711999999999999136e-10 -1.900000000000000105e-04 2.000000000000000000e+02 4.723000000000000337e-09 9.481000000000000657e-10 -1.812000000000000134e-04 2.000000000000000000e+02 4.623000000000000269e-09 9.244999999999999611e-10 -1.727999999999999990e-04 2.000000000000000000e+02 4.520999999999999587e-09 9.002999999999999826e-10 -1.646999999999999918e-04 2.000000000000000000e+02 4.415999999999999640e-09 8.756999999999999849e-10 -1.570999999999999968e-04 2.000000000000000000e+02 4.308999999999999906e-09 8.506000000000000405e-10 -1.498000000000000091e-04 2.000000000000000000e+02 4.199000000000000079e-09 8.250999999999999735e-10 -1.428000000000000016e-04 2.000000000000000000e+02 4.087999999999999945e-09 7.991999999999999905e-10 -1.362000000000000038e-04 2.000000000000000000e+02 3.973000000000000239e-09 7.728999999999999883e-10 -1.299000000000000133e-04 2.000000000000000000e+02 3.857000000000000226e-09 7.462999999999999974e-10 -1.238000000000000005e-04 2.000000000000000000e+02 3.738999999999999599e-09 7.192999999999999873e-10 -1.180999999999999975e-04 2.000000000000000000e+02 3.619000000000000014e-09 6.921000000000000192e-10 -1.125999999999999994e-04 2.000000000000000000e+02 3.498000000000000121e-09 6.648000000000000204e-10 -1.073999999999999949e-04 2.000000000000000000e+02 3.375000000000000029e-09 6.372000000000000330e-10 -1.023999999999999954e-04 2.000000000000000000e+02 3.251000000000000044e-09 6.096000000000000456e-10 -9.767999999999999990e-05 2.000000000000000000e+02 3.126000000000000165e-09 5.819000000000000275e-10 -9.315999999999999574e-05 2.000000000000000000e+02 2.999000000000000087e-09 5.543000000000000401e-10 -8.883999999999999645e-05 2.000000000000000000e+02 2.872999999999999902e-09 5.266999999999999493e-10 -8.472999999999999683e-05 2.000000000000000000e+02 2.745999999999999823e-09 4.993000000000000232e-10 -8.080999999999999371e-05 2.000000000000000000e+02 2.619000000000000158e-09 4.720999999999999517e-10 -7.707000000000000580e-05 2.000000000000000000e+02 2.492000000000000080e-09 4.452000000000000239e-10 -7.351000000000000600e-05 2.000000000000000000e+02 2.365999999999999895e-09 4.187000000000000120e-10 -7.011000000000000467e-05 2.000000000000000000e+02 2.241000000000000016e-09 3.926000000000000194e-10 -6.687000000000000181e-05 2.000000000000000000e+02 2.117000000000000031e-09 3.670000000000000251e-10 -6.378000000000000260e-05 2.000000000000000000e+02 1.993999999999999938e-09 3.420000000000000081e-10 -6.082999999999999866e-05 2.000000000000000000e+02 1.873000000000000046e-09 3.176000000000000199e-10 -5.801999999999999676e-05 2.000000000000000000e+02 1.753999999999999940e-09 2.940000000000000187e-10 -5.534000000000000209e-05 2.000000000000000000e+02 1.637999999999999927e-09 2.711000000000000254e-10 -5.279000000000000109e-05 2.000000000000000000e+02 1.525000000000000007e-09 2.490999999999999980e-10 -5.034999999999999735e-05 2.000000000000000000e+02 1.414999999999999973e-09 2.279999999999999881e-10 -4.802999999999999923e-05 2.000000000000000000e+02 1.308000000000000033e-09 2.077999999999999959e-10 -4.580999999999999677e-05 2.000000000000000000e+02 1.205000000000000078e-09 1.886000000000000001e-10 -4.369999999999999834e-05 2.000000000000000000e+02 1.105999999999999903e-09 1.704000000000000009e-10 -4.169000000000000234e-05 2.000000000000000000e+02 1.010999999999999921e-09 1.533000000000000031e-10 -3.975999999999999879e-05 2.000000000000000000e+02 9.204000000000000476e-10 1.373000000000000067e-10 -3.792999999999999768e-05 2.000000000000000000e+02 8.345999999999999924e-10 1.224000000000000116e-10 -3.618999999999999740e-05 2.000000000000000000e+02 7.535999999999999619e-10 1.085000000000000001e-10 -3.452000000000000153e-05 2.000000000000000000e+02 6.773999999999999562e-10 9.575000000000000539e-11 -3.292999999999999812e-05 2.000000000000000000e+02 6.062000000000000366e-10 8.406000000000000492e-11 -3.142000000000000072e-05 2.000000000000000000e+02 5.399999999999999964e-10 7.343999999999999403e-11 diff --git a/examples/2025-amars/amars-ck.yaml b/examples/2025-amars/amars-ck.yaml deleted file mode 100644 index 611d29c..0000000 --- a/examples/2025-amars/amars-ck.yaml +++ /dev/null @@ -1,147 +0,0 @@ -# Radiation configuration file -# -# Required the fields: -# - species: list of species -# - opacities: list of opacities -# - bands: list of bands -# - : band configuration -# -# Required fields for 'species': -# - name: species name -# - composition: species atomic composition -# -# Required fields for 'opacities': -# - type: choose between the following recognized types -# - 'rfm-lbl': RFM line-by-line table -# - 'rfm-ck': correlated-k table generated from RFM line-by-line -# - 'fourcolumn': four-column opacity table -# -# Optional fields for 'opacities': -# - data: list of opacity data files -# - species: list of species used by the opacity source -# -# Required fields for each band: -# - range: band wavenumber/wavelength range -# If the range is in wavenumber, the unit is cm^{-1} -# If the range is in wavelength, the unit is um -# - opacities: list of opacities used by the band -# - integration: choose between 'wavelength', 'wavenumber', or 'weight' -# If the integration is 'weight', the band is integrated -# using weights from input. Otherwise, the band is integrated -# using the trapezoidal rule along a wavelength -# or wavenumber grid (ww). -# - solver: choose between 'disort' or 'twostr' -# - flags: a string of flags for the radiative transfer solver -# separated by commas -# -# Optional fields for each band: -# - ww: a list of weights (double) for spectral bins in the band - -species: - - name: CO2 - composition: {C: 1, O: 2} - - - name: H2O - composition: {H: 2, O: 1} - - - name: SO2 - composition: {S: 1, O: 2} - - - name: H2SO4 - composition: {H: 2, S: 1, O: 4} - - - name: S8 - composition: {S: 8} - -opacities: - H2SO4: - type: fourcolumn - data: [h2so4.txt] - species: [H2SO4] - nmom: 4 - - S8: - type: fourcolumn - data: [s8_k_fuller.txt] - species: [S8] - nmom: 4 - - CO2: - type: rfm-ck - data: [amars-ck-.nc] - species: [CO2] - - H2O: - type: rfm-ck - data: [amars-ck-.nc] - species: [H2O] - - SO2: - type: rfm-ck - data: [amars-ck-.nc] - species: [SO2] - -bands: [SW, B1, B2, B3, B4, B5, B6, B7, B8] - -SW: - range: [2000., 50000.] - opacities: [H2SO4, S8] - solver: disort - integration: wavenumber - flags: lamber,quiet,onlyfl - -B1: - range: [1., 250.] - opacities: [CO2] - solver: disort - integration: weight - flags: lamber,quiet,onlyfl,planck - -B2: - range: [250., 438.] - opacities: [H2O] - solver: disort - integration: weight - flags: lamber,quiet,onlyfl,planck - -B3: - range: [438., 675.] - opacities: [SO2] - integration: weight - solver: disort - flags: lamber,quiet,onlyfl,planck - -B4: - range: [675., 1062.] - opacities: [CO2] - integration: weight - solver: disort - flags: lamber,quiet,onlyfl,planck - -B5: - range: [1062., 1200.] - opacities: [SO2] - integration: weight - solver: disort - flags: lamber,quiet,onlyfl,planck - -B6: - range: [1200., 1600.] - opacities: [H2O] - integration: weight - solver: disort - flags: lamber,quiet,onlyfl,planck - -B7: - range: [1600., 1900.] - opacities: [SO2] - integration: weight - solver: disort - flags: lamber,quiet,onlyfl,planck - -B8: - range: [1900., 2000.] - opacities: [H2O] - integration: weight - solver: disort - flags: lamber,quiet,onlyfl,planck diff --git a/examples/2025-amars/amars_lw.cpp b/examples/2025-amars/amars_lw.cpp deleted file mode 100644 index bdaf11a..0000000 --- a/examples/2025-amars/amars_lw.cpp +++ /dev/null @@ -1,68 +0,0 @@ -// torch -#include - -// harp -#include -#include -#include - -// unit = [mol/m^3] -torch::Tensor atm_concentration(int ncol, int nlyr, int nspecies) { - auto conc = torch::ones({ncol, nlyr, nspecies}, torch::kFloat64); - return conc; -} - -//! \brief Main function for testing AMARS longwave radiation -/*! - * This is an example program that illustrates how to manually set up - * opacity sources and run a radiative transfer calculation. - * Normally, you would initialize the model from a yaml configuration file. - * However, this example shows the internal mechanisms more clearly. - */ -int main(int argc, char** argv) { - harp::species_names = {"CO2", "H2O"}; - harp::species_weights = {44.0e-3, 18.0e-3}; - - auto op_co2 = harp::OpacityOptionsImpl::create(); - (*op_co2).type("rfm-ck").species_ids({0}).opacity_files({"amarsw-ck-B1.nc"}); - - auto op_h2o = harp::OpacityOptionsImpl::create(); - (*op_h2o).type("rfm-ck").species_ids({1}).opacity_files({"amarsw-ck-B1.nc"}); - - auto lw_op = harp::RadiationBandOptionsImpl::create(); - (*lw_op).name("lw").solver_name("disort").opacities({ - {"CO2", op_co2}, - {"H2O", op_h2o}, - }); - - int nwave = op_co2->query_wavenumber().size(); - int ncol = 1; - int nlyr = 40; - double wmin = 1.; - double wmax = 150.; - - lw_op->nwave(nwave); - lw_op->ncol(ncol); - lw_op->nlyr(nlyr); - - lw_op->disort() = - harp::create_disort_config_lw(wmin, wmax, nwave, ncol, nlyr); - lw_op->weight() = op_co2->query_weight(); - - harp::RadiationBand lw(lw_op); - - auto conc = atm_concentration(ncol, nlyr, harp::species_names.size()); - - std::map bc; - bc["albedo"] = torch::ones({nwave, ncol}, torch::kFloat64); - bc["btemp"] = torch::ones({ncol}, torch::kFloat64) * 300.0; - - std::map atm; - atm["pres"] = torch::ones({ncol, nlyr}, torch::kFloat64) * 10.e5; - atm["temp"] = torch::ones({ncol, nlyr}, torch::kFloat64) * 300.0; - - auto dz = torch::ones({nlyr}, torch::kFloat64); - auto flux = lw->forward(conc, dz, &bc, &atm); - std::cout << "spectrum shape = " << lw->spectrum.sizes() << std::endl; - std::cout << "result = " << flux << std::endl; -} diff --git a/examples/2025-amars/amars_rt.cpp b/examples/2025-amars/amars_rt.cpp deleted file mode 100644 index 2953d19..0000000 --- a/examples/2025-amars/amars_rt.cpp +++ /dev/null @@ -1,224 +0,0 @@ -// C/C++ -#include -#include - -// torch -#include - -// harp -#include -#include -#include -#include -#include -#include -#include - -struct AtmosphericData { - int n_layers; - std::map data; -}; - -AtmosphericData read_rfm_atm(const std::string& filename) { - std::ifstream file(filename); - if (!file.is_open()) { - throw std::runtime_error("Could not open file: " + filename); - } - - AtmosphericData atm_data; - std::string line; - - // Read the number of layers - std::getline(file, line); - std::istringstream iss(line); - iss >> atm_data.n_layers; - - // Read the data sections - while (std::getline(file, line)) { - if (line.empty()) continue; - - // Read the name of the data section - if (line[0] == '*') { - std::string name = line.substr(1); // Remove the '*' character - - // Read the data points - std::getline(file, line); - std::istringstream data_stream(line); - std::vector data_points(atm_data.n_layers); - for (int i = 0; i < atm_data.n_layers; ++i) { - data_stream >> data_points[i]; - } - - // Store the data points in the map - atm_data.data[name] = torch::tensor(data_points, torch::kFloat64); - } - } - - file.close(); - return atm_data; -} - -int main(int argc, char** argv) { - // parameters of the computational grid - int ncol = 1; - int nlyr = 40; - int nstr = 8; - - // parameters of the amars model - double g = 3.711; - double mean_mol_weight = 0.044; // CO2 - double R = 8.314472; - double cp = 844; // J/(kg K) for CO2 - double surf_sw_albedo = 0.3; - // double aero_scale = 1e-6; - double aero_scale = 1; - double sr_sun = 2.92842e-5; // angular size of the sun at mars - double btemp0 = 210; - double solar_temp = 5772; - double lum_scale = 0.7; - - /// ----- read atmosphere data ----- /// - - auto aero_ptx = harp::read_data_tensor("aerosol_output_data.txt"); - auto aero_p = (aero_ptx.select(1, 0) * 1e5); - auto aero_t = aero_ptx.select(1, 1); - auto aero_x = aero_ptx.narrow(1, 2, 2); - - // unit = [pa] - auto new_P = harp::read_data_tensor("pVals.txt").squeeze(1); - new_P *= 100.0; // convert mbar to Pa - - // unit = [K] - auto new_T = harp::read_data_tensor("TVals.txt").squeeze(1); - auto dT_dt = torch::zeros_like(new_T); - - // unit = [mol/mol] - auto new_X = harp::interpn({new_P.log()}, {aero_p.log()}, aero_x); - - AtmosphericData atm_data = read_rfm_atm("rfm.atm"); - auto pre = atm_data.data["PRE [mb]"] * 100.0; - auto tem = atm_data.data["TEM [K]"]; - - // unit = [mol/m^3] - auto conc = torch::zeros({ncol, nlyr, 5}, torch::kFloat64); - - conc.select(-1, 0) = (atm_data.data["CO2 [ppmv]"] * 1e-6) * (pre / (R * tem)); - conc.select(-1, 1) = (atm_data.data["H2O [ppmv]"] * 1e-6) * (pre / (R * tem)); - conc.select(-1, 2) = (atm_data.data["SO2 [ppmv]"] * 1e-6) * (pre / (R * tem)); - - // aerosols - conc.narrow(-1, 3, new_X.size(-1)) = - aero_scale * new_X * new_P.unsqueeze(1) / (R * new_T.unsqueeze(1)); - - // unit = [kg/m^3] - auto new_rho = (new_P * mean_mol_weight) / (R * new_T); - - // unit = [m] - auto dz = harp::calc_dz_hypsometric(new_P, new_T, - torch::tensor({mean_mol_weight * g / R})); - - /// ----- done read atmosphere data ----- /// - - // configure input data for each radiation band - std::map atm, bc; - atm["pres"] = new_P.unsqueeze(0).expand({ncol, nlyr}); - atm["temp"] = new_T.unsqueeze(0).expand({ncol, nlyr}); - atm["rho"] = new_rho.unsqueeze(0).expand({ncol, nlyr}); - - // read radiation configuration from yaml file - auto op = harp::RadiationOptions::from_yaml("amars-ck.yaml"); - for (auto& [name, band] : op.bands()) { - // query weights from opacity, only valid for longwave - // shortwave values are defined separately - band.ww() = band.query_weights(); - int nwave = name == "SW" ? 500 : band.ww().size(); - - auto wmin = band.disort().wave_lower()[0]; - auto wmax = band.disort().wave_upper()[0]; - - harp::disort_config(&band.disort(), nwave, ncol, nlyr, nstr); - - if (name == "SW") { // shortwave - band.ww().resize(nwave); - for (int i = 0; i < nwave; ++i) { - band.ww()[i] = (wmax - wmin) * i / (nwave - 1) + wmin; - } - auto wave = torch::tensor(band.ww(), torch::kFloat64); - bc[name + "/fbeam"] = - lum_scale * sr_sun * harp::bbflux_wavenumber(wave, solar_temp, ncol); - bc[name + "/albedo"] = - surf_sw_albedo * torch::ones({nwave, ncol}, torch::kFloat64); - bc[name + "/umu0"] = 0.707 * torch::ones({ncol}, torch::kFloat64); - } else { // longwave - band.disort().wave_lower(std::vector(nwave, wmin)); - band.disort().wave_upper(std::vector(nwave, wmax)); - bc[name + "/albedo"] = 0.0 * torch::ones({nwave, ncol}, torch::kFloat64); - } - } - bc["btemp"] = btemp0 * torch::ones({ncol}, torch::kFloat64); - bc["ttemp"] = torch::zeros({ncol}, torch::kFloat64); - - // print radiation options and construct radiation model - harp::Radiation rad(op); - auto [netflux, dnflux, upflux] = rad->forward(conc, dz, &bc, &atm); - - int t_lim = 10000; - int print_freq = 500; - double tstep = 86400 / 4.; - double cSurf = - 200000; // thermal inertia of the surface, assuming half ocean half land - double current_time = 0; - - for (int t_ind = 0; t_ind < t_lim; ++t_ind) { - auto surf_forcing = dnflux - 5.67e-8 * bc["btemp"].pow(4); - bc["btemp"] += surf_forcing * (tstep / cSurf); - - dT_dt = torch::zeros_like(atm["temp"]); - auto dT_dt = -1. / (atm["rho"] * cp) * - (netflux.narrow(-1, 1, nlyr) - netflux.narrow(-1, 0, nlyr)) / - dz; - - atm["temp"] += dT_dt * tstep; - atm["temp"].clamp_(20, 1000); - - conc.select(-1, 0) = (atm_data.data["CO2 [ppmv]"] * 1e-6) * - (atm["pres"] / (R * atm["temp"])); - conc.select(-1, 1) = (atm_data.data["H2O [ppmv]"] * 1e-6) * - (atm["pres"] / (R * atm["temp"])); - conc.select(-1, 2) = (atm_data.data["SO2 [ppmv]"] * 1e-6) * - (atm["pres"] / (R * atm["temp"])); - - // aerosols - conc.narrow(-1, 3, new_X.size(-1)) = aero_scale * new_X.unsqueeze(0) * - atm["pres"].unsqueeze(-1) / - (R * atm["temp"].unsqueeze(-1)); - - // unit = [kg/m^3] - atm["rho"] = (atm["pres"] * mean_mol_weight) / (R * atm["temp"]); - - // unit = [m] - dz = harp::calc_dz_hypsometric(atm["pres"], atm["temp"], - torch::tensor({mean_mol_weight * g / R})); - - if (t_ind % print_freq == 0) { - std::ostringstream filename; - filename << "tp_result" << t_ind << ".txt"; - - // Open the file and write the data - std::ofstream outputFile3(filename.str()); - outputFile3 << "#p[Pa] T[K] netF[w/m^2] dT/dt [K/s]" << std::endl; - for (int k = 0; k < nlyr; ++k) { - outputFile3 << atm["pres"][0][k].item() << " " - << atm["temp"][0][k].item() << " " - << atm["temp"][0][k].item() << " " - << dT_dt[0][k].item() << std::endl; - } - outputFile3.close(); - } - - auto result = rad->forward(conc, dz, &bc, &atm); - netflux = std::get<0>(result); - dnflux = std::get<1>(result); - upflux = std::get<2>(result); - } -} diff --git a/examples/2025-amars/amars_rt.py b/examples/2025-amars/amars_rt.py deleted file mode 100644 index fc31854..0000000 --- a/examples/2025-amars/amars_rt.py +++ /dev/null @@ -1,117 +0,0 @@ -#! /user/bin/env python3 -import torch -import numpy as np -import pyharp as harp -from torch import tensor, logspace, zeros, ones -from pyharp import ( - interpn, - constants, - calc_dz_hypsometric, - bbflux_wavenumber, - RadiationOptions, - Radiation, - read_rfm_atm, -) - -if __name__ == "__main__": - ncol = 1 - nlyr = 80 - nstr = 4 - - surf_sw_albedo = 0.3 - sr_sun = 2.92842e-5 - btemp0 = 210 - ttemp0 = 100 - solar_temp = 5772 - lum_scale = 0.7 - grav = 3.711 - mean_mol_weight = 0.044 - - pres = logspace(np.log10(500.0), np.log10(0.01), nlyr) - pres *= 100.0 - temp = 200.0 * ones((nlyr,), dtype=torch.float64) - - # mole fractions - xfrac = zeros((ncol, nlyr, 5), dtype=torch.float64) - - # molecules - rfm_atm = read_rfm_atm("rfm.atm") - rfm_pre = rfm_atm["PRE"] * 100.0 - rfm_tem = rfm_atm["TEM"] - - xfrac[:, :, 0] = interpn( - [pres.log()], [rfm_pre.log()], rfm_atm["CO2"].unsqueeze(-1) * 1.0e-6 - ).squeeze(-1) - xfrac[:, :, 1] = interpn( - [pres.log()], [rfm_pre.log()], rfm_atm["H2O"].unsqueeze(-1) * 1.0e-6 - ).squeeze(-1) - xfrac[:, :, 2] = interpn( - [pres.log()], [rfm_pre.log()], rfm_atm["SO2"].unsqueeze(-1) * 1.0e-6 - ).squeeze(-1) - - # aerosols - aero_ptx = tensor(np.genfromtxt("aerosol_output_data.txt")) - aero_p = aero_ptx[:, 0] * 1.0e5 - aero_t = aero_ptx[:, 1] - aero_x = aero_ptx[:, 2:] - - xfrac[:, :, 3:] = interpn([pres.log()], [aero_p.log()], aero_x) - atm = { - "pres": pres.unsqueeze(0).expand(ncol, nlyr), - "temp": temp.unsqueeze(0).expand(ncol, nlyr), - } - bc = {} - - # layer thickness - dz = calc_dz_hypsometric( - pres, temp, tensor(mean_mol_weight * grav / constants.Rgas) - ) - - rad_op = RadiationOptions.from_yaml("amars-ck.yaml") - - # configure bands - for band in rad_op.bands(): - name = band.name() - # Query wavenumber and weight from the opacity sources - op = list(band.opacities().values())[0] if band.opacities() else None - if op and name != "SW": - band.wavenumber(list(op.query_wavenumber())) - band.weight(list(op.query_weight())) - nwave = len(band.wavenumber()) if name != "SW" else 200 - - wmin = band.disort().wave_lower()[0] - wmax = band.disort().wave_upper()[0] - - band.disort().accur(1.0e-12) - - if name == "SW": # shortwave - band.wavenumber(list(np.linspace(wmin, wmax, nwave))) - wave = tensor(band.wavenumber(), dtype=torch.float64) - bc[name + "/fbeam"] = ( - lum_scale * sr_sun * bbflux_wavenumber(wave, solar_temp) - ) - bc[name + "/albedo"] = surf_sw_albedo * ones( - (nwave, ncol), dtype=torch.float64 - ) - bc[name + "/umu0"] = 0.707 * ones((ncol,), dtype=torch.float64) - else: # longwave - band.disort().wave_lower([wmin] * nwave) - band.disort().wave_upper([wmax] * nwave) - bc[name + "/albedo"] = zeros((nwave, ncol), dtype=torch.float64) - bc[name + "/temis"] = ones((nwave, ncol), dtype=torch.float64) - - bc["btemp"] = btemp0 * ones((ncol,), dtype=torch.float64) - bc["ttemp"] = ttemp0 * ones((ncol,), dtype=torch.float64) - - # construct radiation model - print("radiation options:\n", rad_op) - rad = Radiation(rad_op) - - # run RT - conc = xfrac.clone() - conc *= atm["pres"].unsqueeze(-1) / (constants.Rgas * atm["temp"].unsqueeze(-1)) - netflux, dnflux, upflux = rad.forward(conc, dz, bc, atm) - - print("net flux = ", netflux) - print("downward flux = ", dnflux) - print("upward flux = ", upflux) diff --git a/examples/2025-amars/amars_rt_cli.cpp b/examples/2025-amars/amars_rt_cli.cpp deleted file mode 100644 index df168fc..0000000 --- a/examples/2025-amars/amars_rt_cli.cpp +++ /dev/null @@ -1,176 +0,0 @@ -// C/C++ -#include -#include - -// torch -#include - -// harp -#include -#include -#include -#include -#include -#include -#include - -struct AtmosphericData { - int n_layers; - std::map data; -}; - -AtmosphericData read_rfm_atm(const std::string& filename) { - std::ifstream file(filename); - if (!file.is_open()) { - throw std::runtime_error("Could not open file: " + filename); - } - - AtmosphericData atm_data; - std::string line; - - // Read the number of layers - std::getline(file, line); - std::istringstream iss(line); - iss >> atm_data.n_layers; - - // Read the data sections - while (std::getline(file, line)) { - if (line.empty()) continue; - - // Read the name of the data section - if (line[0] == '*') { - std::string name = line.substr(1); // Remove the '*' character - - // Read the data points - std::getline(file, line); - std::istringstream data_stream(line); - std::vector data_points(atm_data.n_layers); - for (int i = 0; i < atm_data.n_layers; ++i) { - data_stream >> data_points[i]; - } - - // Store the data points in the map - atm_data.data[name] = torch::tensor(data_points, torch::kFloat64); - } - } - - file.close(); - return atm_data; -} - -int main(int argc, char** argv) { - // parameters of the computational grid - int ncol = 1; - int nlyr = 40; - int nstr = 8; - - // parameters of the amars model - double g = 3.711; - double mean_mol_weight = 0.044; // CO2 - double R = 8.314472; - double cp = 844; // J/(kg K) for CO2 - double surf_sw_albedo = 0.3; - // double aero_scale = 1e-6; - double aero_scale = 1; - double sr_sun = 2.92842e-5; // angular size of the sun at mars - double btemp = 210; - double solar_temp = 5772; - double lum_scale = 0.7; - - /// ----- read atmosphere data ----- /// - - auto aero_ptx = harp::read_data_tensor("aerosol_output_data.txt"); - auto aero_p = (aero_ptx.select(1, 0) * 1e5); - auto aero_t = aero_ptx.select(1, 1); - auto aero_x = aero_ptx.narrow(1, 2, 2); - std::cout << "aero_p = " << aero_p << std::endl; - std::cout << "aero_t = " << aero_t << std::endl; - std::cout << "aero_x = " << aero_x.sizes() << std::endl; - - // unit = [pa] - auto new_P = harp::read_data_tensor("pVals.txt").squeeze(1); - new_P *= 100.0; // convert mbar to Pa - std::cout << "new_P = " << new_P << std::endl; - - // unit = [K] - auto new_T = harp::read_data_tensor("TVals.txt").squeeze(1); - std::cout << "new_T = " << new_T << std::endl; - - // unit = [mol/mol] - auto new_X = harp::interpn({new_P.log()}, {aero_p.log()}, aero_x); - std::cout << "new_X = " << new_X << std::endl; - - AtmosphericData atm_data = read_rfm_atm("rfm.atm"); - auto pre = atm_data.data["PRE [mb]"] * 100.0; - auto tem = atm_data.data["TEM [K]"]; - - // unit = [mol/m^3] - auto conc = torch::zeros({ncol, nlyr, 5}, torch::kFloat64); - - conc.select(-1, 0) = (atm_data.data["CO2 [ppmv]"] * 1e-6) * (pre / (R * tem)); - conc.select(-1, 1) = (atm_data.data["H2O [ppmv]"] * 1e-6) * (pre / (R * tem)); - conc.select(-1, 2) = (atm_data.data["SO2 [ppmv]"] * 1e-6) * (pre / (R * tem)); - - // aerosols - conc.narrow(-1, 3, new_X.size(-1)) = - aero_scale * new_X * new_P.unsqueeze(1) / (R * new_T.unsqueeze(1)); - std::cout << "conc = " << conc << std::endl; - - // unit = [kg/m^3] - auto new_rho = (new_P * mean_mol_weight) / (R * new_T); - std::cout << "new_rho = " << new_rho << std::endl; - - // unit = [m] - auto dz = harp::calc_dz_hypsometric(new_P, new_T, - torch::tensor({mean_mol_weight * g / R})); - std::cout << "dz = " << dz << std::endl; - - /// ----- done read atmosphere data ----- /// - - // configure input data for each radiation band - std::map atm, bc; - atm["pres"] = new_P.unsqueeze(0).expand({ncol, nlyr}); - atm["temp"] = new_T.unsqueeze(0).expand({ncol, nlyr}); - - // read radiation configuration from yaml file - auto op = harp::RadiationOptions::from_yaml("amars-ck.yaml"); - for (auto& [name, band] : op.bands()) { - // query weights from opacity, only valid for longwave - // shortwave values are defined separately - band.ww() = band.query_weights(); - int nwave = name == "SW" ? 500 : band.ww().size(); - - auto wmin = band.disort().wave_lower()[0]; - auto wmax = band.disort().wave_upper()[0]; - - harp::disort_config(&band.disort(), nwave, ncol, nlyr, nstr); - std::cout << "flags = " << band.disort().flags() << std::endl; - - if (name == "SW") { // shortwave - band.ww().resize(nwave); - for (int i = 0; i < nwave; ++i) { - band.ww()[i] = (wmax - wmin) * i / (nwave - 1) + wmin; - } - auto wave = torch::tensor(band.ww(), torch::kFloat64); - bc[name + "/fbeam"] = - lum_scale * sr_sun * harp::bbflux_wavenumber(wave, solar_temp, ncol); - bc[name + "/albedo"] = - surf_sw_albedo * torch::ones({nwave, ncol}, torch::kFloat64); - bc[name + "/umu0"] = 0.707 * torch::ones({nwave, ncol}, torch::kFloat64); - } else { // longwave - band.disort().wave_lower(std::vector(nwave, wmin)); - band.disort().wave_upper(std::vector(nwave, wmax)); - bc[name + "/albedo"] = 0.0 * torch::ones({nwave, ncol}, torch::kFloat64); - } - } - bc["btemp"] = btemp * torch::ones({ncol}, torch::kFloat64); - bc["ttemp"] = torch::zeros({ncol}, torch::kFloat64); - - // print radiation options and construct radiation model - std::cout << "rad op = " << fmt::format("{}", op) << std::endl; - harp::Radiation rad(op); - auto [netflux, dnflux, upflux] = rad->forward(conc, dz, &bc, &atm); - std::cout << "net flux = " << netflux << std::endl; - std::cout << "downward flux = " << dnflux << std::endl; - std::cout << "upward flux = " << upflux << std::endl; -} diff --git a/examples/2025-amars/amars_rt_cm.cpp b/examples/2025-amars/amars_rt_cm.cpp deleted file mode 100644 index 99eb854..0000000 --- a/examples/2025-amars/amars_rt_cm.cpp +++ /dev/null @@ -1,710 +0,0 @@ -// harp -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -std::vector> read_4width_array_from_file( - std::string fpath) { - std::ifstream file(fpath); - std::vector> array_to_get; - std::string line; - - std::getline(file, line); // skip the first line - while (std::getline(file, line)) { - std::istringstream iss(line); - std::vector row; - double value; - while (iss >> value) { - row.push_back(value); - } - if (row.size() == 4) { - array_to_get.push_back(row); - } else { - std::cerr << "Invalid line format: " << line << std::endl; - } - } - file.close(); - - return array_to_get; -} - -std::vector read_values_from_file(const std::string& filename) { - std::ifstream file(filename); - if (!file.is_open()) { - throw std::runtime_error("Could not open file: " + filename); - } - - std::vector values; - double value; - while (file >> value) { - values.push_back(value); - } - - file.close(); - return values; -} - -struct AtmosphericData { - int n_layers; - std::map> data; -}; - -AtmosphericData read_atmospheric_data(const std::string& filename) { - std::ifstream file(filename); - if (!file.is_open()) { - throw std::runtime_error("Could not open file: " + filename); - } - - AtmosphericData atm_data; - std::string line; - - // Read the number of layers - std::getline(file, line); - std::istringstream iss(line); - iss >> atm_data.n_layers; - - // Read the data sections - while (std::getline(file, line)) { - if (line.empty()) continue; - - // Read the name of the data section - if (line[0] == '*') { - std::string name = line.substr(1); // Remove the '*' character - - // Read the data points - std::getline(file, line); - std::istringstream data_stream(line); - std::vector data_points(atm_data.n_layers); - for (int i = 0; i < atm_data.n_layers; ++i) { - data_stream >> data_points[i]; - } - - // Store the data points in the map - atm_data.data[name] = data_points; - } - } - - file.close(); - return atm_data; -} - -torch::Tensor calc_dz_from_file(int nlyr, torch::Tensor prop, - AtmosphericData atm_data) { - auto dz = torch::ones({nlyr, 1}, - prop.options()); // prop: (nwave, ncol, nlyr, nprop) - double z_i = 0; - double z_i_plus_1 = 0; - - dz[0][0] = atm_data.data["HGT [km]"][0] * 1000; - for (int i = 1; i < nlyr - 1; ++i) { - z_i = atm_data.data["HGT [km]"][i]; - z_i_plus_1 = atm_data.data["HGT [km]"][i + 1]; - dz[i][0] *= (z_i_plus_1 - z_i) * 1000; - } - dz[nlyr - 1][0] *= (z_i_plus_1 - z_i) * 1000; - return dz; -} - -// unit = [mol/m^3] -torch::Tensor read_atm_concentration(int ncol, int nlyr, int nspecies, - AtmosphericData atm_data) { - auto conc = torch::ones({ncol, nlyr, nspecies}, torch::kFloat64); - double pre = 0; - double tem = 0; - double R = 8.314472; - for (int i = 0; i < nlyr; ++i) { - pre = atm_data.data["PRE [mb]"][i] * 100.0; - tem = atm_data.data["TEM [K]"][i]; - conc[0][i][0] = (atm_data.data["CO2 [ppmv]"][i] * 1e-6) * (pre / (R * tem)); - conc[0][i][1] = atm_data.data["H2O [ppmv]"][i] * 1e-6 * (pre / (R * tem)); - conc[0][i][2] = atm_data.data["SO2 [ppmv]"][i] * 1e-6 * (pre / (R * tem)); - } - return conc; -} - -torch::Tensor calc_flux_1band_init(int ncol, int nspecies, double wmin, - double wmax, AtmosphericData atm_data, - std::string filename, int idx, - double btemp) { - harp::AttenuatorOptions op; - op.species_names({"CO2", "H2O", "SO2"}); - op.species_weights({44.0e-3, 18.0e-3, 64.0e-3}); - - op.species_ids({0}).opacity_files({filename}); - harp::RFM co2(op); - - op.species_ids({1}).opacity_files({filename}); - harp::RFM h2o(op); - - op.species_ids({2}).opacity_files({filename}); - harp::RFM so2(op); - - int nwave = co2->kdata.size(0); - int nlyr = atm_data.n_layers; - - auto conc = read_atm_concentration(ncol, nlyr, nspecies, atm_data); - - disort::Disort disort(disort_options_lw(wmin, wmax, nwave, ncol, nlyr)); - - std::map kwargs; - kwargs["pres"] = torch::ones({ncol, nlyr}, torch::kFloat64); - kwargs["temp"] = torch::ones({ncol, nlyr}, torch::kFloat64); - for (int i = 0; i < nlyr; ++i) { - kwargs["pres"][0][i] = atm_data.data["PRE [mb]"][i] * 100.0; - kwargs["temp"][0][i] = atm_data.data["TEM [K]"][i]; - } - - torch::Tensor prop; - if (idx == 0) prop = co2->forward(conc, kwargs); - if (idx == 1) prop = h2o->forward(conc, kwargs); - if (idx == 2) prop = so2->forward(conc, kwargs); - auto dz = calc_dz_from_file(nlyr, prop, atm_data); - - prop *= dz; - - std::map bc; - bc["albedo"] = torch::ones({nwave, ncol}, torch::kFloat64) * - 0.0; // leave emissivity at 1 - bc["btemp"] = torch::ones({nwave, ncol}, torch::kFloat64) * btemp; - - auto temf = harp::layer2level(kwargs["temp"], harp::Layer2LevelOptions()); - auto flux = disort->forward(prop, &bc, temf); - auto weights = harp::read_weights_rfm(filename); - - return (flux * weights.view({-1, 1, 1, 1})).sum(0); -} - -torch::Tensor calc_flux_1band_loop(int ncol, int nspecies, double wmin, - double wmax, AtmosphericData atm_data, - std::string filename, int idx, double btemp, - torch::Tensor new_T, torch::Tensor new_p, - torch::Tensor dz) { - harp::AttenuatorOptions op; - op.species_names({"CO2", "H2O", "SO2"}); - op.species_weights({44.0e-3, 18.0e-3, 64.0e-3}); - - op.species_ids({0}).opacity_files({filename}); - harp::RFM co2(op); - - op.species_ids({1}).opacity_files({filename}); - harp::RFM h2o(op); - - op.species_ids({2}).opacity_files({filename}); - harp::RFM so2(op); - - int nwave = co2->kdata.size(0); - int nlyr = atm_data.n_layers; - - auto conc = torch::ones({ncol, nlyr, nspecies}, torch::kFloat64); - double pre = 0; - double tem = 0; - double R = 8.314472; - for (int i = 0; i < nlyr; ++i) { - pre = new_p[i]; - tem = new_T[i]; - conc[0][i][0] = (atm_data.data["CO2 [ppmv]"][i] * 1e-6) * (pre / (R * tem)); - conc[0][i][1] = atm_data.data["H2O [ppmv]"][i] * 1e-6 * (pre / (R * tem)); - conc[0][i][2] = atm_data.data["SO2 [ppmv]"][i] * 1e-6 * (pre / (R * tem)); - } - - disort::Disort disort(disort_options_lw(wmin, wmax, nwave, ncol, nlyr)); - - std::map kwargs; - kwargs["pres"] = torch::ones({ncol, nlyr}, torch::kFloat64); - kwargs["temp"] = torch::ones({ncol, nlyr}, torch::kFloat64); - kwargs["pres"][0] = new_p; - kwargs["temp"][0] = new_T; - - torch::Tensor prop; - if (idx == 0) prop = co2->forward(conc, kwargs); - if (idx == 1) prop = h2o->forward(conc, kwargs); - if (idx == 2) prop = so2->forward(conc, kwargs); - - prop *= dz; - - std::map bc; - bc["albedo"] = torch::ones({nwave, ncol}, torch::kFloat64) * - 0.0; // leave emissivity at 1 - bc["btemp"] = torch::ones({nwave, ncol}, torch::kFloat64) * btemp; - - auto temf = harp::layer2level(kwargs["temp"], harp::Layer2LevelOptions()); - auto flux = disort->forward(prop, &bc, temf); - auto weights = harp::read_weights_rfm(filename); - - return (flux * weights.view({-1, 1, 1, 1})).sum(0); -} - -std::tuple calc_flux_1band_loop_with_prop( - int ncol, int nspecies, double wmin, double wmax, AtmosphericData atm_data, - std::string filename, int idx, double btemp, std::vector new_T, - std::vector new_p, torch::Tensor dz, int t_ind, int print_freq) { - harp::AttenuatorOptions op; - op.species_names({"CO2", "H2O", "SO2"}); - op.species_weights({44.0e-3, 18.0e-3, 64.0e-3}); - - op.species_ids({0}).opacity_files({filename}); - harp::RFM co2(op); - - op.species_ids({1}).opacity_files({filename}); - harp::RFM h2o(op); - - op.species_ids({2}).opacity_files({filename}); - harp::RFM so2(op); - - int nwave = co2->kdata.size(0); - int nlyr = atm_data.n_layers; - - auto conc = torch::ones({ncol, nlyr, nspecies}, torch::kFloat64); - double pre = 0; - double tem = 0; - double R = 8.314472; - for (int i = 0; i < nlyr; ++i) { - pre = new_p[i]; - tem = new_T[i]; - conc[0][i][0] = atm_data.data["CO2 [ppmv]"][i] * 1e-6 * (pre / (R * tem)); - conc[0][i][1] = atm_data.data["H2O [ppmv]"][i] * 1e-6 * (pre / (R * tem)); - conc[0][i][2] = atm_data.data["SO2 [ppmv]"][i] * 1e-6 * (pre / (R * tem)); - } - - disort::Disort disort(disort_options_lw(wmin, wmax, nwave, ncol, nlyr)); - - std::map kwargs; - kwargs["pres"] = torch::ones({ncol, nlyr}, torch::kFloat64); - kwargs["temp"] = torch::ones({ncol, nlyr}, torch::kFloat64); - for (int i = 0; i < nlyr; ++i) { - kwargs["pres"][0][i] = new_p[i]; - kwargs["temp"][0][i] = new_T[i]; - } - - torch::Tensor prop; - if (idx == 0) prop = co2->forward(conc, kwargs); - if (idx == 1) prop = h2o->forward(conc, kwargs); - if (idx == 2) prop = so2->forward(conc, kwargs); - - prop *= dz; - - std::map bc; - bc["albedo"] = torch::ones({nwave, ncol}, torch::kFloat64) * - 0.0; // leave emissivity at 1 - bc["btemp"] = torch::ones({nwave, ncol}, torch::kFloat64) * btemp; - - auto temf = harp::layer2level(kwargs["temp"], harp::Layer2LevelOptions()); - auto flux = disort->forward(prop, &bc, temf); - auto weights = harp::read_weights_rfm(filename); - - auto flux_result = (flux * weights.view({-1, 1, 1, 1})).sum(0); - - if (t_ind % print_freq == 0) { - std::ostringstream filename3; - filename3 << "conc_result_lw" << t_ind + 1 << ".txt"; - std::ofstream outputFile6(filename3.str()); - outputFile6 << "#p[Pa] conc_co2 conc_h2o conc_so2" << std::endl; - for (int k = 0; k < nlyr; ++k) { - outputFile6 << new_p[k] << " "; - outputFile6 << conc[0][k][0].item() << " "; - outputFile6 << conc[0][k][1].item() << " "; - outputFile6 << conc[0][k][2].item() << std::endl; - } - outputFile6 << std::endl; - outputFile6.close(); - } - - return std::make_tuple(flux_result, prop); -} - -double calculate_dynamic_timestep(const std::vector& new_T, - const std::vector& dT_ds, int nlyr, - double safety_factor) { - double min_time_to_zero = std::numeric_limits::max(); - bool found_negative = false; - - for (int k = 0; k < nlyr; ++k) { - if (dT_ds[k] < 0) { - double time_to_zero = new_T[k] / std::abs(dT_ds[k]); - if (time_to_zero < min_time_to_zero) { - min_time_to_zero = time_to_zero; - } - found_negative = true; - } - } - - if (!found_negative) { - return 1e6; // Return a large default timestep if no negative values are - // found - } - - double timestep = min_time_to_zero / safety_factor; - return (timestep > 1e6) ? 1e6 : timestep; -} - -//// ------- shortwave ------- //// -void calc_shortwave(double wmin, double wmax, int nwave, int nlyr) { - double solar_temp = 5772; - double lum_scale = 0.7; - int ncol = 1; - int nwave = wave.size(0); - - // shortwave grid from 0.2um to 5um (2000 cm^-1 to 50000 cm^-1) - auto wave = torch::linspace(2000, 50000, nwave, torch::kFloat64); - - RadiationBandOptions rt_sw_op; - rt_sw_op.name() = "sw"; - rt_sw_op.solver_name() = "disort"; - rt_sw_op.opacity() = { - {"h2so4", - op.species_ids({0}).opacity_files({"h2so4.txt"}).type("h2so4_simple")}, - {"s8", op.species_ids({1}) - .opacity_files({"s8_k_fuller.txt"}) - .type("fourcolumn")}}; - rt_sw_op.disort() = disort(disort_flux_sw(nwave, ncol, nlyr)); - - RadiationBand rt_sw(rt_sw_op); - std::map other; - other["wavenumber"] = wave; - - std::map bc; - bc["fbeam"] = lum_scale * sr_sun * bb_flux(wave, solar_temp, ncol); - bc["umu0"] = 0.707 * torch::ones({nwave, ncol}, torch::kFloat64); - bc["albedo"] = surf_sw_albedo * torch::ones({nwave, ncol}, torch::kFloat64); - - // shortwave flux at each wavenumber/level - return rt_sw->forward(conc, dz, bc, other); -} - -//// ------- longwave ------- //// -void calc_longwave_1band(std::string filename, double wmin, double wmax, - int nlyr) { - RadiationBandOptions rt_lw_op; - rt_lw_op.name() = "lw"; - rt_lw_op.solver_name() = "disort"; - rt_lw_op.opacity() = { - {"CO2", op.species_ids({0}).opacity_files({filename}).type("rfm")}, - {"H2O", op.species_ids({1}).opacity_files({filename}).type("rfm")}, - {"SO2", op.species_ids({2}).opacity_files({filename}).type("rfm")}}; - rt_lw_op.disort() = disort_flux_lw(wmin, wmax, nwave, ncol, nlyr); - - RadiationBand rt_lw(rt_lw_op); - - // longwave flux at each wavenumber/level - return rt_lw->forward(conc, dz, bc, other); -} - -void calc_longwave(std::string filename, int nlyr) { - auto tot_flux = - calc_longwave_1band(1., 250., atm_data, "amars-ck-B1.nc", 0, btemp); - tot_flux += - calc_flux_1band_init(250., 438., atm_data, "amars-ck-B2.nc", 1, btemp); - tot_flux += - calc_flux_1band_init(438., 675., atm_data, "amars-ck-B3.nc", 2, btemp); - tot_flux += - calc_flux_1band_init(675., 1062., atm_data, "amars-ck-B4.nc", 0, btemp); - tot_flux += - calc_flux_1band_init(1062., 1200., atm_data, "amars-ck-B5.nc", 2, btemp); - tot_flux += - calc_flux_1band_init(1200., 1600., atm_data, "amars-ck-B6.nc", 0, btemp); - tot_flux += - calc_flux_1band_init(1600., 1900., atm_data, "amars-ck-B7.nc", 2, btemp); - tot_flux += - calc_flux_1band_init(1900., 2000., atm_data, "amars-ck-B8.nc", 0, btemp); -} - -int main(int argc, char** argv) { - // int nwave = 48; - int nwave = 500; // 50 bins gets you within ~1 W/m^2 fldn at TOA, but we want - // to be sure to resolve spectral info - // int nwave = 15000; //essentially the exact value of the integral over the - // chosen wavelength bounds - int ncol = 1; - int nlyr = 40; // 3 layers is 1 W/m^2 away from the exact value of fldn_surf - // when using 200 layers. however, we want some more layers to - // resolve heating - int nspecies = 2; - double g = 3.711; - double mean_mol_weight = 0.044; // CO2 - double R = 8.314472; - double cp = 844; // J/(kg K) for CO2 - double surf_sw_albedo = 0.3; - // double aero_scale = 1e-6; - double aero_scale = 1; - double sr_sun = 2.92842e-5; // angular size of the sun at mars - - //// ------- species ------- //// - harp::AttenuatorOptions op; - op.species_names({"H2SO4", "S8", "CO2", "H2O", "SO2"}); - op.species_weights({98.e-3, 256.e-3, 44.0e-3, 18.0e-3, 64.0e-3}); - - // read in the atmos output, and extract pressure and mixing ratios - // std::vector> aero_mr_p = - // read_4width_array_from_file("aerosol_output_data.txt"); - - auto aero_mr_p = harp::read_data_tensor("aerosol_output_data.txt"); - aeros_mr_p.select(1, 0) *= 1e5; // convert bar to Pa - - /*std::vector p(aero_mr_p.size()); - std::vector T(aero_mr_p.size()); - std::vector> mr(nspecies, - std::vector(aero_mr_p.size())); - for (size_t i = 0; i < aero_mr_p.size(); ++i) { - p[i] = aero_mr_p[i][0] * 1e5; // convert bar to Pa - T[i] = aero_mr_p[i][1]; - for (int j = 0; j < nspecies; ++j) { - mr[j][i] = aero_mr_p[i][j + 2]; - } - }*/ - - // unit = [pa] - auto new_P = harp::read_data_tensor("pVals.txt"); - new_P *= 100.0; // convert mbar to Pa - - // unit = [K] - auto new_T = harp::read_data_tensor("TVals.txt"); - - // unit = [mol/mol] - auto new_X = harp::interpn({new_P.log()}, {aero_mr_p.select(1, 0).log()}, - aero_mr_p.narrow(1, 2, nspecies)); - - // unit = [mol/m^3] - auto conc = aero_scale * new_X * new_P / (R * new_T); - - // unit = [kg/m^3] - auto new_rho = (new_P * mean_mol_weight) / (R * new_T); - - auto dz = calc_dz_hypsometric(new_p, new_T, torch::tensor({g / R})); - - auto total_flux_sw = harp::cal_total_flux(flux_sw, wave, "wave"); - auto surf_flux_sw = harp::cal_surface_flux(total_flux_sw); - auto toa_flux_sw = harp::cal_toa_flux(total_flux_sw); - - // std::cout << "tot_flux_down_surf: " << tot_flux_down_surf << " W/m^2" - // << std::endl; - // std::cout << "tot_flux_down_toa: " << tot_flux_down_toa << " W/m^2" - // << std::endl; - - //// ------- longwave ------- //// - - double btemp = 210; - AtmosphericData atm_data = read_atmospheric_data("rfm.atm"); - int nlyr_lw = atm_data.n_layers; - int nspecies_lw = 3; - // you must pass the index of the species of the associated ck band - // CO2 is idx 0, H2O is idx 1, SO2 idx is 2 - // std::cout << "tot_flux = " << tot_flux << std::endl; - - // calculate heating rates and write to file - // unit = [K/s] - auto dTdt = torch::zeros({nlyr}, torch::kFloat64); - - double df = 0; - double df_iplus1 = 0; - std::ofstream outputFile("dT_ds.txt"); - outputFile << "#p[Pa] dT_ds[K/s]" << std::endl; - for (int k = 0; k < nlyr; ++k) { - // add SW fluxes up/down to the LW fluxes up/down. 0 is up, 1 is down - df = integrated_flux[k][0] - integrated_flux[k][1] + - tot_flux[0][k][0].item() - tot_flux[0][k][1].item(); - df_iplus1 = integrated_flux[k + 1][0] - integrated_flux[k + 1][1] + - tot_flux[0][k + 1][0].item() - - tot_flux[0][k + 1][1].item(); - dT_ds[k] = - -(1 / (new_rho[k] * cp)) * (df_iplus1 - df) / dz[k][0].item(); - outputFile << new_p[k] << " " << dT_ds[k] << std::endl; - } - outputFile.close(); - - std::ostringstream filename; - filename << "tp_result" << 0 << ".txt"; - - // Open the file and write the data - std::ofstream outputFile_init(filename.str()); - outputFile_init << "#p[Pa] T[K] dT/dt [K/s]" << std::endl; - for (int k = 0; k < nlyr; ++k) { - outputFile_init << new_p[k] << " " << new_T[k] << " " << dT_ds[k] - << std::endl; - } - outputFile_init.close(); - - // CALC THE RAD EQUIL - int t_lim = 50000; // 10000 = 6 yr with 6hr timesteps - int print_freq = 500; - double tstep = 86400 / 1; // step is 6hrs - double cSurf = - 200000; // thermal inertia of the surface, assuming half ocean half land - double current_time = 0; - std::ofstream outputFile2("btemp_t.txt"); - outputFile2 << "#t[s] btemp[K]" << std::endl; - outputFile2 << current_time << " " << btemp << std::endl; - - // start the time loop - for (int t_ind = 0; t_ind < t_lim; ++t_ind) { - // double tstep = calculate_dynamic_timestep(new_T, dT_ds, nlyr, 1000); - // std::cout << "tstep = " << tstep << std::endl; - // double tstep = 86400; - - current_time += tstep; - // sw + lw flux down - surf emission - double surf_forcing = (1 - surf_sw_albedo) * integrated_flux[0][1] + - tot_flux[0][0][1].item() - - 5.67e-8 * std::pow(btemp, 4); - btemp += surf_forcing * (tstep / cSurf); - - for (int k = 0; k < nlyr; ++k) { - new_T[k] += dT_ds[k] * tstep; - if (new_T[k] < 20) new_T[k] = 20; - - // PRESSURE GRID IS FIXED - conc[0][k][0] = - (aero_scale * new_mr[1][k] * new_p[k]) / - (R * new_T[k]); // s8 comes second in the file that we read - // in. but we need it to be index 0 in conc - // bc of how it was initialized above - conc[0][k][1] = - (aero_scale * new_mr[0][k] * new_p[k]) / - (R * new_T[k]); // h2so4 comes first in the file that we read in. but - // it needs to be index 1 in conc. - new_rho[k] = (new_p[k] * mean_mol_weight) / (R * new_T[k]); - } - - auto prop1 = s8->forward(conc, kwargs); - auto prop2 = h2so4->forward(conc, kwargs); - auto prop = prop1 + prop2; - - auto dz = calc_dz(nlyr, prop, new_p, new_rho, g); - prop *= dz; - - // mean single scattering albedo - prop.select(3, 1) /= prop.select(3, 0); - - auto result = disort->forward(prop, &bc); - - auto [integrated_flux, tot_flux_down_surf, tot_flux_down_toa] = - integrate_result(result, wave, nlyr, nwave); - - if (t_ind % print_freq == 0) { - std::ostringstream filename3; - filename3 << "conc_result_sw" << t_ind + 1 << ".txt"; - std::ofstream outputFile6(filename3.str()); - outputFile6 << "#p[Pa] conc_s8 conc_h2so4" << std::endl; - for (int k = 0; k < nlyr; ++k) { - outputFile6 << new_p[k] << " "; - outputFile6 << conc[0][k][0].item() << " "; - outputFile6 << conc[0][k][1].item() << " "; - } - outputFile6 << std::endl; - outputFile6.close(); - } - - std::vector prop_results; - torch::Tensor tot_flux; - torch::Tensor flux_result, prop_result; - - std::tie(flux_result, prop_result) = calc_flux_1band_loop_with_prop( - ncol, nspecies_lw, 1., 250., atm_data, "amars-ck-B1.nc", 0, btemp, - new_T, new_p, dz, t_ind, print_freq); - tot_flux = flux_result; - prop_results.push_back(prop_result); - - std::tie(flux_result, prop_result) = calc_flux_1band_loop_with_prop( - ncol, nspecies_lw, 250., 438., atm_data, "amars-ck-B2.nc", 1, btemp, - new_T, new_p, dz, t_ind, print_freq); - tot_flux += flux_result; - prop_results.push_back(prop_result); - - std::tie(flux_result, prop_result) = calc_flux_1band_loop_with_prop( - ncol, nspecies_lw, 438., 675., atm_data, "amars-ck-B3.nc", 2, btemp, - new_T, new_p, dz, t_ind, print_freq); - tot_flux += flux_result; - prop_results.push_back(prop_result); - - std::tie(flux_result, prop_result) = calc_flux_1band_loop_with_prop( - ncol, nspecies_lw, 675., 1062., atm_data, "amars-ck-B4.nc", 0, btemp, - new_T, new_p, dz, t_ind, print_freq); - tot_flux += flux_result; - prop_results.push_back(prop_result); - - std::tie(flux_result, prop_result) = calc_flux_1band_loop_with_prop( - ncol, nspecies_lw, 1062., 1200., atm_data, "amars-ck-B5.nc", 2, btemp, - new_T, new_p, dz, t_ind, print_freq); - tot_flux += flux_result; - prop_results.push_back(prop_result); - - std::tie(flux_result, prop_result) = calc_flux_1band_loop_with_prop( - ncol, nspecies_lw, 1200., 1600., atm_data, "amars-ck-B6.nc", 0, btemp, - new_T, new_p, dz, t_ind, print_freq); - tot_flux += flux_result; - prop_results.push_back(prop_result); - - std::tie(flux_result, prop_result) = calc_flux_1band_loop_with_prop( - ncol, nspecies_lw, 1600., 1900., atm_data, "amars-ck-B7.nc", 2, btemp, - new_T, new_p, dz, t_ind, print_freq); - tot_flux += flux_result; - prop_results.push_back(prop_result); - - std::tie(flux_result, prop_result) = calc_flux_1band_loop_with_prop( - ncol, nspecies_lw, 1900., 2000., atm_data, "amars-ck-B8.nc", 0, btemp, - new_T, new_p, dz, t_ind, print_freq); - tot_flux += flux_result; - prop_results.push_back(prop_result); - - // calculate heating rates - double df = 0; - double df_iplus1 = 0; - for (int k = 0; k < nlyr; ++k) { - // add SW fluxes up/down to the LW fluxes up/down. 0 is up, 1 is down - df = integrated_flux[k][0] - integrated_flux[k][1] + - tot_flux[0][k][0].item() - tot_flux[0][k][1].item(); - df_iplus1 = integrated_flux[k + 1][0] - integrated_flux[k + 1][1] + - tot_flux[0][k + 1][0].item() - - tot_flux[0][k + 1][1].item(); - dT_ds[k] = - -(1 / (new_rho[k] * cp)) * (df_iplus1 - df) / dz[k][0].item(); - } - outputFile2 << current_time << " " << btemp << std::endl; - - if (t_ind % print_freq == 0) { - std::ostringstream filename; - filename << "tp_result" << t_ind + 1 << ".txt"; - - // Open the file and write the data - std::ofstream outputFile3(filename.str()); - outputFile3 << "#p[Pa] T[K] dT/dt [K/s]" << std::endl; - for (int k = 0; k < nlyr; ++k) { - outputFile3 << new_p[k] << " " << new_T[k] << " " << dT_ds[k] - << std::endl; - } - outputFile3.close(); - - std::ostringstream filename2; - filename2 << "tau_result" << t_ind + 1 << ".txt"; - std::ofstream outputFile5(filename2.str()); - outputFile5 << "#p[Pa] tau_sw ssa_sw tau_b1 tau_b2 tau_b3 tau_b4 tau_b5 " - "tau_b6 tau_b7 tau_b8" - << std::endl; - for (int k = 0; k < nlyr; ++k) { - outputFile5 << new_p[k] << " "; - outputFile5 << prop.select(3, 0)[0][0][k].item() << " "; - outputFile5 << prop.select(3, 1)[0][0][k].item() << " "; - for (int i = 0; i < 8; ++i) { - outputFile5 << prop_results[i].select(3, 0)[0][0][k].item() - << " "; - } - outputFile5 << std::endl; - } - outputFile5.close(); - } - } - outputFile2.close(); - - std::ofstream outputFile4("final_tp_result.txt"); - outputFile4 << "#p[Pa] T[K]" << std::endl; - for (int k = 0; k < nlyr; ++k) { - outputFile4 << new_p[k] << " " << new_T[k] << std::endl; - } - outputFile4.close(); -} diff --git a/examples/2025-amars/amars_rt_rk3.cpp b/examples/2025-amars/amars_rt_rk3.cpp deleted file mode 100644 index 5208541..0000000 --- a/examples/2025-amars/amars_rt_rk3.cpp +++ /dev/null @@ -1,245 +0,0 @@ -// C/C++ -#include -#include - -// torch -#include - -// harp -#include -#include -#include -#include -#include -#include - -struct AtmosphericData { - int n_layers; - std::map data; -}; - -AtmosphericData read_rfm_atm(const std::string& filename) { - std::ifstream file(filename); - if (!file.is_open()) { - throw std::runtime_error("Could not open file: " + filename); - } - - AtmosphericData atm_data; - std::string line; - - // Read the number of layers - std::getline(file, line); - std::istringstream iss(line); - iss >> atm_data.n_layers; - - // Read the data sections - while (std::getline(file, line)) { - if (line.empty()) continue; - - // Read the name of the data section - if (line[0] == '*') { - std::string name = line.substr(1); // Remove the '*' character - - // Read the data points - std::getline(file, line); - std::istringstream data_stream(line); - std::vector data_points(atm_data.n_layers); - for (int i = 0; i < atm_data.n_layers; ++i) { - data_stream >> data_points[i]; - } - - // Store the data points in the map - atm_data.data[name] = torch::tensor(data_points, torch::kFloat64); - } - } - - file.close(); - return atm_data; -} - -int main(int argc, char** argv) { - // parameters of the computational grid - int ncol = 1; - int nlyr = 80; - int nstr = 4; - - // parameters of the amars model - double surf_sw_albedo = 0.3; - double sr_sun = 2.92842e-5; // angular size of the sun at mars - double btemp0 = 210; - double ttemp0 = 100; - double solar_temp = 5772; - double lum_scale = 0.7; - - /// ----- read atmosphere data ----- /// - auto aero_ptx = harp::read_data_tensor("aerosol_output_data.txt"); - auto aero_p = (aero_ptx.select(1, 0) * 1e5); // bar to pa - auto aero_t = aero_ptx.select(1, 1); - auto aero_x = aero_ptx.narrow(1, 2, 2); - - // unit = [pa] - // log pressure grid from 500 mbar to 1 mbar - auto new_P = - torch::logspace(log10(500.0), log10(1.0), nlyr, 10, torch::kFloat64); - new_P *= 100.0; // convert mbar to Pa - - // unit = [K] - // isothermal temperature profile at 200 K - auto new_T = 200. * torch::ones({nlyr}, torch::kFloat64); - - // unit = [mol/mol] - auto new_X = harp::interpn({new_P.log()}, {aero_p.log()}, aero_x); - - AtmosphericData atm_data = read_rfm_atm("rfm.atm"); - auto pre = atm_data.data["PRE [mb]"] * 100.0; - auto tem = atm_data.data["TEM [K]"]; - - // unit = [mol/mol] - // mole fraction - auto xfrac = torch::zeros({ncol, nlyr, 5}, torch::kFloat64); - xfrac.select(-1, 0) = - harp::interpn({new_P.log()}, {pre.log()}, - atm_data.data["CO2 [ppmv]"].unsqueeze(-1) * 1e-6) - .squeeze(-1); - - xfrac.select(-1, 1) = - harp::interpn({new_P.log()}, {pre.log()}, - atm_data.data["H2O [ppmv]"].unsqueeze(-1) * 1e-6) - .squeeze(-1); - - xfrac.select(-1, 2) = - harp::interpn({new_P.log()}, {pre.log()}, - atm_data.data["SO2 [ppmv]"].unsqueeze(-1) * 1e-6) - .squeeze(-1); - - // aerosols - xfrac.narrow(-1, 3, new_X.size(-1)) = new_X; - - /// ----- done read atmosphere data ----- /// - std::map atm, bc; - // set up model atmosphere - atm["pres"] = new_P.unsqueeze(0).expand({ncol, nlyr}); - atm["temp"] = new_T.unsqueeze(0).expand({ncol, nlyr}); - - // read radiation configuration from yaml file - auto rad_op = harp::RadiationOptions::from_yaml("amars-ck.yaml"); - for (auto& [name, band] : rad_op.bands()) { - // query weights from opacity, only valid for longwave - // shortwave values are defined separately - band.ww() = band.query_weights(); - int nwave = name == "SW" ? 200 : band.ww().size(); - - auto wmin = band.disort().wave_lower()[0]; - auto wmax = band.disort().wave_upper()[0]; - band.disort().accur(1.e-12); - - harp::disort_config(&band.disort(), nwave, ncol, nlyr, nstr); - - if (name == "SW") { // shortwave - band.ww().resize(nwave); - for (int i = 0; i < nwave; ++i) { - band.ww()[i] = (wmax - wmin) * i / (nwave - 1) + wmin; - } - auto wave = torch::tensor(band.ww(), torch::kFloat64); - bc[name + "/fbeam"] = - lum_scale * sr_sun * harp::bbflux_wavenumber(wave, solar_temp, ncol); - bc[name + "/albedo"] = - surf_sw_albedo * torch::ones({nwave, ncol}, torch::kFloat64); - bc[name + "/umu0"] = 0.707 * torch::ones({ncol}, torch::kFloat64); - } else { // longwave - band.disort().wave_lower(std::vector(nwave, wmin)); - band.disort().wave_upper(std::vector(nwave, wmax)); - bc[name + "/albedo"] = 0.0 * torch::ones({nwave, ncol}, torch::kFloat64); - bc[name + "/temis"] = 1.0 * torch::ones({nwave, ncol}, torch::kFloat64); - } - } - bc["btemp"] = btemp0 * torch::ones({ncol}, torch::kFloat64); - bc["ttemp"] = ttemp0 * torch::ones({ncol}, torch::kFloat64); - - // parameters of the amars model - harp::RadiationModelOptions model_op; - model_op.ncol(ncol); - model_op.nlyr(nlyr); - model_op.grav(3.711); - model_op.mean_mol_weight(0.044); // CO2 - model_op.cp(844); // J/(kg K) for CO2 - model_op.aero_scale(1.); - model_op.cSurf(200000); // J/(m^2 K) thermal intertia of the surface - model_op.intg(harp::IntegratorOptions().type("rk2")); - model_op.kappa(2.e-2); // m^2/s thermal diffusivity - model_op.rad(rad_op); - - harp::RadiationModel model(model_op); - - int t_lim = 10000; - double tstep = 86400 / 4.; - int print_freq = 500; - - for (int t_ind = 0; t_ind < t_lim; ++t_ind) { - for (int stage = 0; stage < model->pintg->stages.size(); ++stage) { - model->forward(xfrac, atm, bc, tstep, stage); - } - - if (t_ind % print_freq == 0) { - std::ostringstream filename; - filename << "tp_result" << t_ind << ".txt"; - - // Open the file and write the data - std::ofstream outputFile3(filename.str()); - - // opacity of SW band - auto prop = harp::shared["radiation/SW/opacity"]; - - outputFile3 << "#p[Pa] T[K] netF[w/m^2] dT/dt[K/s] SW_tau1 SW_tau2 " - << "B1 B2 B3 B4 B5 B6 B7 B8" << std::endl; - for (int k = 0; k < nlyr; ++k) { - outputFile3 << atm["pres"][0][k].item() << " " - << atm["temp"][0][k].item() << " " - << harp::shared["result/netflux"][0][k].item() - << " " << harp::shared["result/dT_atm"][0][k].item() - << " " - << harp::shared["radiation/SW/opacity"] - .mean(0)[0][k][0] - .item() - << " " - << harp::shared["radiation/SW/opacity"] - .mean(0)[0][k][1] - .item() - << " " - << harp::shared["radiation/B1/opacity"] - .mean(0)[0][k][0] - .item() - << " " - << harp::shared["radiation/B2/opacity"] - .mean(0)[0][k][0] - .item() - << " " - << harp::shared["radiation/B3/opacity"] - .mean(0)[0][k][0] - .item() - << " " - << harp::shared["radiation/B4/opacity"] - .mean(0)[0][k][0] - .item() - << " " - << harp::shared["radiation/B5/opacity"] - .mean(0)[0][k][0] - .item() - << " " - << harp::shared["radiation/B6/opacity"] - .mean(0)[0][k][0] - .item() - << " " - << harp::shared["radiation/B7/opacity"] - .mean(0)[0][k][0] - .item() - << " " - << harp::shared["radiation/B8/opacity"] - .mean(0)[0][k][0] - .item() - << " " << std::endl; - } - outputFile3.close(); - } - } -} diff --git a/examples/2025-amars/amars_sw.py b/examples/2025-amars/amars_sw.py deleted file mode 100644 index 0577053..0000000 --- a/examples/2025-amars/amars_sw.py +++ /dev/null @@ -1,252 +0,0 @@ -// C/C++ -#include -#include -#include -#include -#include -#include - -// torch -#include - -// base -#include - -// harp -#include -#include - -#include -#include -#include -#include -#include -#include -#include - -void interpolate_mixing_ratios(const std::vector& dense_grid, - const std::vector& mixing_ratios, - const std::vector& thinned_grid, - std::vector& interpolated_ratios) { - // Ensure the output vector is the correct size - interpolated_ratios.resize(thinned_grid.size()); - - // Perform 1D interpolation for each point in the thinned grid - for (size_t i = 0; i < thinned_grid.size(); ++i) { - interpolated_ratios[i] = interp1(thinned_grid[i], mixing_ratios.data(), - dense_grid.data(), dense_grid.size()); - } -} - -// unit = [cm^-1] -std::vector short_wavenumber_grid(int nwave) { - int wmin = 2000; // 5 um - int wmax = 50000; // 200 nm this gets us within 2 W/m^2 of the correct 410 - // value. this is where rayleigh attenuates all intensity - std::vector wave_vec(nwave); - double dw = (wmax - wmin) / (nwave - 1); - for (int i = 0; i < nwave; ++i) { - wave_vec[i] = wmin + i * dw; - } - return wave_vec; -} - -// unit = [w/(m^2 cm^-1)] -torch::Tensor short_toa_flux(int nwave, int ncol) { - return torch::ones({nwave, ncol}, torch::kFloat64); -} - -// unit = [w/(m^2 cm^-1)] -torch::Tensor bb_toa_flux(torch::Tensor wave, int ncol, double temp, - double fscale) { - double c1 = 1.19144e-5 * 1e-3; - double c2 = 1.4388; - double sr_sun = 2.92842e-5; // angular size of the sun at mars - - int nwave = wave.size(0); - torch::Tensor bb_flux = torch::ones({nwave, ncol}, torch::kFloat64); - for (int i = 0; i < nwave; ++i) { - for (int j = 0; j < ncol; ++j) { - bb_flux[i][j] = fscale * sr_sun * c1 * pow(wave[i], 3) / - (exp(c2 * wave[i] / temp) - 1); - } - } - return bb_flux; -} - -std::vector> read_4width_array_from_file( - std::string fpath) { - std::ifstream file(fpath); - std::vector> array_to_get; - std::string line; - - std::getline(file, line); // skip the first line - while (std::getline(file, line)) { - std::istringstream iss(line); - std::vector row; - double value; - while (iss >> value) { - row.push_back(value); - } - if (row.size() == 4) { - array_to_get.push_back(row); - } else { - std::cerr << "Invalid line format: " << line << std::endl; - } - } - file.close(); - - return array_to_get; -} - -std::tuple, std::vector, - std::vector>> -regrid_ptx(int nlyr, int nspecies, std::vector p, std::vector T, - std::vector> mr) { - double p_min = *std::min_element(p.begin(), p.end()); - double p_max = *std::max_element(p.begin(), p.end()); - double p_step = (p_max - p_min) / (nlyr - 1); - - double T_min = *std::min_element(T.begin(), T.end()); - double T_max = *std::max_element(T.begin(), T.end()); - double T_step = (T_max - T_min) / (nlyr - 1); - - std::vector new_p(nlyr); - std::vector new_T(nlyr); - for (int i = 0; i < nlyr; ++i) { - new_p[nlyr - 1 - i] = p_min + i * p_step; - new_T[nlyr - 1 - i] = T_min + i * T_step; - } - - std::vector> new_mr(nspecies, std::vector(nlyr)); - for (int j = 0; j < nspecies; ++j) { - interpolate_mixing_ratios(p, mr[j], new_p, new_mr[j]); - } - - return std::make_tuple(new_p, new_T, new_mr); -} - -int main(int argc, char** argv) { - // int nwave = 48; - int nwave = 500; // 50 bins gets you within ~1 W/m^2 fldn at TOA, but we want - // to be sure to resolve spectral info - // int nwave = 15000; //essentially the exact value of the integral over the - // chosen wavelength bounds - int ncol = 1; - int nlyr = 40; // 3 layers is 1 W/m^2 away from the exact value of fldn_surf - // when using 200 layers. however, we want some more layers to - // resolve heating - int nspecies = 2; - double g = 3.711; - double mean_mol_weight = 0.044; // CO2 - double R = 8.314472; - double cp = 844; // J/(kg K) for CO2 - double solar_temp = 5772; - double lum_scale = 0.7; - - auto conc = torch::zeros({ncol, nlyr, nspecies}, torch::kFloat64); - - // read in the atmos output, and extract pressure and mixing ratios - std::vector> aero_mr_p = - read_4width_array_from_file("aerosol_output_data.txt"); - std::vector P(aero_mr_p.size()); - std::vector T(aero_mr_p.size()); - std::vector> X(nspecies, - std::vector(aero_mr_p.size())); - for (size_t i = 0; i < aero_mr_p.size(); ++i) { - P[i] = aero_mr_p[i][0] * 1e5; // convert bar to Pa - T[i] = aero_mr_p[i][1]; - for (int j = 0; j < nspecies; ++j) { - X[j][i] = aero_mr_p[i][j + 2]; - } - } - - auto [P, T, X] = regrid_ptx(nlyr, nspecies, P, T, X); - std::vector rho(nlyr); - - for (int k = 0; k < nlyr; ++k) { - conc[0][k][0] = (X[1][k] * P[k]) / - (R * T[k]); // s8 comes second in the file that we read - // in. but we need it to be index 0 in conc - // bc of how it was initialized above - conc[0][k][1] = - (X[0][k] * P[k]) / - (R * T[k]); // h2so4 comes first in the file that we read in. but - // it needs to be index 1 in conc. - rho[k] = (P[k] * mean_mol_weight) / (R * T[k]); - } - - harp::species_names = {"S8", "H2SO4"}; - harp::species_weights = {256.e-3, 98.e-3}; - - auto op_s8 = harp::OpacityOptionsImpl::create(); - auto op_h2so4 = harp::OpacityOptionsImpl::create(); - - (*op_s8).type("fourcolumn").species_ids({0}).opacity_files({"s8_k_fuller.txt"}); - - (*op_h2so4).type("fourcolumn").species_ids({1}).opacity_files({"h2so4.txt"}); - - auto sw_op = harp::RadiationBandOptionsImpl::create(); - (*sw_op).name("sw").solver_name("disort").opacities({ - {"S8", op_s8}, - {"H2SO4", op_h2so4}, - }); - - sw_op->nwave(nwave); - sw_op->ncol(ncol); - sw_op->nlyr(nlyr); - - sw_op->disort() = harp::create_disort_config_sw(nwave, ncol, nlyr); - sw_op->wavenumber() = short_wavenumber_grid(nwave); - - harp::RadiationBand sw(sw_op); - - std::map bc; - bc["fbeam"] = bb_toa_flux(wave, ncol, solar_temp, lum_scale); - bc["umu0"] = 1.0 * torch::ones({nwave, ncol}, torch::kFloat64); - bc["albedo"] = 1.0 * torch::ones({nwave, ncol}, torch::kFloat64); - - std::map atm; - atm["pres"] = torch::tensor(new_p, torch::kFloat64).unsqueeze(0); - atm["temp"] = torch::tensor(new_T, torch::kFloat64).unsqueeze(0); - - auto result = sw->forward(conc, dz, &bc, &atm); - - auto [integrated_flux, tot_flux_down_surf, tot_flux_down_toa] = - integrate_result(result, wave, nlyr, nwave); - - std::cout << "tot_flux_down_surf: " << tot_flux_down_surf << " W/m^2" - << std::endl; - std::cout << "tot_flux_down_toa: " << tot_flux_down_toa << " W/m^2" - << std::endl; - - // calculate heating rates and write to file - double df = 0; - double df_iplus1 = 0; - std::vector dTdt(nlyr); - std::ofstream outputFile("dTdt.txt"); - outputFile << "#p[Pa] dTdt[K/s]" << std::endl; - for (int k = 0; k < nlyr; ++k) { - df = integrated_flux[k][0] - integrated_flux[k][1]; - df_iplus1 = integrated_flux[k + 1][0] - integrated_flux[k + 1][1]; - dTdt[k] = - -(1 / (rho[k] * cp)) * (df_iplus1 - df) / dz[k][0].item(); - outputFile << new_p[k] << " " << dTdt[k] << std::endl; - } - outputFile.close(); - - // make sure the mixing ratio interpolation is working fine - std::ofstream outputFile2("mix.txt"); - outputFile2 << "#p_new[Pa] mr(new)(0) mr(new)(1) p_atmos[Pa] mr(atmos)(0) " - "mr(atmos)(1)" - << std::endl; - for (int k = 0; k < nlyr; ++k) { - outputFile2 << new_p[k] << " " << X[0][k] << " " << X[1][k] << " " - << p[k] << " " << mr[0][k] << " " << mr[1][k] << std::endl; - } - for (int k = nlyr; k < 200; ++k) { - outputFile2 << 0 << " " << 0 << " " << 0 << " " << p[k] << " " << mr[0][k] - << " " << mr[1][k] << std::endl; - } - outputFile2.close(); -} diff --git a/examples/2025-amars/amarsw-ck.yaml b/examples/2025-amars/amarsw-ck.yaml deleted file mode 100644 index b571f0a..0000000 --- a/examples/2025-amars/amarsw-ck.yaml +++ /dev/null @@ -1,90 +0,0 @@ -opacity-sources: - - name: CO2 - class: HitranCK - data: amarsw-ck-B1.nc - dependent-species: [vapor.dry] - - - name: H2O - class: HitranCK - data: amarsw-ck-B2.nc - dependent-species: [vapor.H2O] - - - name: CO2 - class: HitranCK - data: amarsw-ck-B3.nc - dependent-species: [vapor.dry] - - - name: H2O - class: HitranCK - data: amarsw-ck-B4.nc - dependent-species: [vapor.H2O] - - - name: CO2 - class: HitranCK - data: amarsw-ck-B5.nc - dependent-species: [vapor.dry] - -bands: [B1, B2, B3, B4, B5] - -#CO2 CIA -B1: - units: cm-1 - grid-type: cktable - wavenumber-range: [1., 150.] - opacity: [CO2] - rt-solver: Disort - flags: [thermal_emission] - -#H2O -B2: - units: cm-1 - grid-type: cktable - wavenumber-range: [150., 500.] - opacity: [H2O] - rt-solver: Disort - flags: [thermal_emission] - -#CO2 + CIA -B3: - units: cm-1 - grid-type: cktable - wavenumber-range: [500., 1450.] - opacity: [CO2] - rt-solver: Disort - flags: [thermal_emission] - -#H2O -B4: - units: cm-1 - grid-type: cktable - wavenumber-range: [1450., 1850.] - opacity: [H2O] - rt-solver: Disort - flags: [thermal_emission] - -#CO2 -B5: - units: cm-1 - grid-type: cktable - wavenumber-range: [1850., 2000.] - opacity: [CO2] - rt-solver: Disort - flags: [thermal_emission] - -Disort-flags: - ibcnd: false - usrtau: false - usrang: false - lamber: true - onlyfl: true - spher: false - intensity_correction: true - old_intensity_correction: false - general_source: false - output_uum: false - quiet: true - print-input: false - print-fluxes: false - print-intensity: false - print-transmissivity: false - print-phase-function: true diff --git a/examples/2025-amars/amarsw-lbl.yaml b/examples/2025-amars/amarsw-lbl.yaml deleted file mode 100644 index de77afc..0000000 --- a/examples/2025-amars/amarsw-lbl.yaml +++ /dev/null @@ -1,95 +0,0 @@ -opacity-sources: - - name: CO2 - class: Hitran - data: amarsw-lbl-B1.nc - dependent-species: [vapor.dry] - - - name: H2O - class: Hitran - data: amarsw-lbl-B2.nc - dependent-species: [vapor.H2O] - - - name: CO2 - class: Hitran - data: amarsw-lbl-B3.nc - dependent-species: [vapor.dry] - - - name: H2O - class: Hitran - data: amarsw-lbl-B4.nc - dependent-species: [vapor.H2O] - - - name: CO2 - class: Hitran - data: amarsw-lbl-B5.nc - dependent-species: [vapor.dry] - -bands: [B1, B2, B3, B4, B5] - -#CO2 CIA -B1: - units: cm-1 - grid-type: regular - wavenumber-range: [1., 150.] - resolution: 0.1 - opacity: [CO2] - rt-solver: Disort - flags: [thermal_emission] - -#H2O -B2: - units: cm-1 - grid-type: regular - wavenumber-range: [150., 500.] - resolution: 0.1 - opacity: [H2O] - rt-solver: Disort - flags: [thermal_emission] - -#CO2 + CIA -B3: - units: cm-1 - grid-type: regular - wavenumber-range: [500., 1450.] - resolution: 0.1 - opacity: [CO2] - rt-solver: Disort - flags: [thermal_emission] - -#H2O -B4: - units: cm-1 - grid-type: regular - wavenumber-range: [1450., 1850.] - resolution: 0.1 - opacity: [H2O] - rt-solver: Disort - flags: [thermal_emission] - -#CO2 -B5: - units: cm-1 - grid-type: regular - wavenumber-range: [1850., 2000.] - resolution: 0.1 - opacity: [CO2] - rt-solver: Disort - flags: [thermal_emission] - -Disort-flags: - ibcnd: false - usrtau: false - usrang: false - lamber: true - onlyfl: true - spher: false - intensity_correction: true - old_intensity_correction: false - general_source: false - output_uum: false - quiet: true - print-input: false - print-fluxes: false - print-intensity: false - print-transmissivity: false - print-phase-function: true diff --git a/examples/2025-amars/pVals.txt b/examples/2025-amars/pVals.txt deleted file mode 100644 index 07872ea..0000000 --- a/examples/2025-amars/pVals.txt +++ /dev/null @@ -1,40 +0,0 @@ -5.000000000000000568e+02 -4.948410268773598091e+02 -4.893343943975348225e+02 -4.833171296553780394e+02 -4.767489454840567191e+02 -4.695878678579084067e+02 -4.617905091652220335e+02 -4.533124389095848983e+02 -4.441086695085879796e+02 -4.341342768161232470e+02 -4.233451770638372977e+02 -4.116990841699270618e+02 -3.991566739277247393e+02 -3.856829846524165646e+02 -3.712490876579888663e+02 -3.558340656275945548e+02 -3.394273424191782169e+02 -3.220314132778462408e+02 -3.036650274505955167e+02 -2.843668706937816637e+02 -2.641997735663443905e+02 -2.432554170658096098e+02 -2.216593981416969825e+02 -1.995763299297577760e+02 -1.772143728529226223e+02 -1.548282496925444889e+02 -1.304131499985158769e+02 -1.096760317644856997e+02 -9.058443075716989767e+01 -7.333779446612462039e+01 -5.808046468824829134e+01 -4.459774689460296315e+01 -3.362420335340597433e+01 -2.472912978536872330e+01 -1.771027668245888620e+01 -1.232887730675849269e+01 -8.180472952810108112e+00 -5.392845876861435350e+00 -3.497085397460474532e+00 -2.231580531608204243e+00 diff --git a/examples/2025-amars/rfm.atm b/examples/2025-amars/rfm.atm deleted file mode 100644 index e342c9d..0000000 --- a/examples/2025-amars/rfm.atm +++ /dev/null @@ -1,14 +0,0 @@ -40 -*HGT [km] -0 0.10435512 0.21960023 0.34687177 0.48742473 0.64264512 0.81406355 1.0033704 1.2124323 1.443311 1.698283 1.9798626 2.2908266 2.6342412 3.0134929 3.4323214 3.8948568 4.4056602 4.9697686 5.5927445 6.2807311 7.0405127 7.8795814 8.8062112 9.8295396 10.959657 12.207709 13.586001 15.108125 16.78909 18.645473 20.695579 22.959624 25.459933 28.221162 31.27054 34.638137 38.357159 42.46428 47 -*PRE [mb] -500 494.84103 489.33439 483.31713 476.74895 469.58787 461.79051 453.31244 444.10867 434.13428 423.34518 411.69908 399.15667 385.68298 371.24909 355.83407 339.42734 322.03141 303.66503 284.36687 264.19977 243.25542 221.6594 199.57633 177.21437 154.82825 130.41315 109.67603 90.584431 73.337794 58.080465 44.597747 33.624203 24.72913 17.710277 12.328877 8.180473 5.3928459 3.4970854 2.2315805 -*TEM [K] -210 209.54084 209.03376 208.47376 207.85533 207.17236 206.41812 205.58517 204.6653 203.64943 202.52755 201.2886 199.92036 198.40934 196.74063 194.89779 192.86263 190.6151 188.13302 185.39192 182.36478 179.02174 175.32984 171.25267 166.75003 161.77751 156.28608 156.28608 156.28608 156.28608 156.28608 156.28608 158.55013 161.05043 163.81166 166.86104 170.22864 180.6419 192.14184 204.84186 -*H2O [ppmv] -14.035208 13.300899 12.527148 11.720344 10.884451 10.024612 9.1472271 8.2599816 7.3718227 6.4928508 5.6341145 4.8072905 4.0242393 3.2964382 2.6343086 2.0464792 1.5390495 1.1149461 0.77348093 0.51022604 0.31729663 0.18407744 0.098341614 0.047603418 0.020454657 0.0075989265 0.0023998515 0.0023998515 0.0023998515 0.0023998515 0.0023998515 0.0023998515 0.0023998515 0.0023998515 0.0023998515 0.0023998515 0.0023998515 0.0023998515 0.0023998515 0.0023998515 -*CO2 [ppmv] -914272.88 916022.24 918533.6 921247 924170.89 927312.09 930675.13 934261.51 938068.8 942089.65 946310.6 950710.88 955261.12 959922.1 964643.72 969364.38 974011.03 978500.33 982741.25 986639.73 990105.57 993061.79 995455.76 997270.67 998534.27 999320.89 999737.75 999737.75 999737.75 999737.75 999737.75 999737.75 999737.75 999737.75 999737.75 999737.75 999737.75 999737.75 999737.75 999737.75 -*SO2 [ppmv] -85713.083 83964.461 81453.871 78741.281 75818.224 72677.883 69315.72 65730.23 61923.826 57903.861 53683.77 49284.312 44734.854 40074.604 35353.648 30633.578 25987.432 21498.559 17257.975 13359.76 9894.111 6938.027 4544.1404 2729.287 1465.7142 679.10339 262.24384 262.24384 262.24384 262.24384 262.24384 262.24384 262.24384 262.24384 262.24384 262.24384 262.24384 262.24384 262.24384 262.24384 -*END diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 82cfdb0..0102b3e 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -1,5 +1,14 @@ -# root project if (${CMAKE_SOURCE_DIR} STREQUAL ${PROJECT_SOURCE_DIR}) - add_subdirectory(2025-amars) + setup_problem(jupiter_1d_rt) + + configure_file( + ${CMAKE_CURRENT_SOURCE_DIR}/jupiter_1d_rt.yaml + ${CMAKE_BINARY_DIR}/bin/jupiter_1d_rt.yaml + COPYONLY) + + #configure_file( + # ${CMAKE_CURRENT_SOURCE_DIR}/h2_0p9_he_0p1_ch4_0p004_h2o_0p002_nh3_0p0003_xsection_120_730K_1_100bar_20_2500cm1.nc + # ${CMAKE_BINARY_DIR}/bin/h2_0p9_he_0p1_ch4_0p004_h2o_0p002_nh3_0p0003_xsection_120_730K_1_100bar_20_2500cm1.nc + # COPYONLY) endif() diff --git a/examples/example_sonora_2020_flux.py b/examples/example_sonora_2020_flux.py index 19e6c34..2be42b3 100644 --- a/examples/example_sonora_2020_flux.py +++ b/examples/example_sonora_2020_flux.py @@ -70,10 +70,10 @@ def init_radiation(config_file: str) -> Radiation: band.disort().accur(1.0e-4) data = [wmin] * ng - band.disort().wave_lower([x for col in zip(*data) for x in col]) + band.set_wave_lower([x for col in zip(*data) for x in col]) data = [wmax] * ng - band.disort().wave_upper([x for col in zip(*data) for x in col]) + band.set_wave_upper([x for col in zip(*data) for x in col]) wave_lower = np.array(band.disort().wave_lower()) wave_upper = np.array(band.disort().wave_upper()) diff --git a/examples/jupiter_1d_rt.cpp b/examples/jupiter_1d_rt.cpp new file mode 100644 index 0000000..be95a39 --- /dev/null +++ b/examples/jupiter_1d_rt.cpp @@ -0,0 +1,318 @@ +// C/C++ +#include +#include +#include +#include +#include +#include +#include +#include + +// external +#include + +// torch +#include + +// harp +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +std::vector tensor_to_vector(torch::Tensor tensor) { + tensor = tensor.to(torch::kCPU).to(torch::kFloat64).contiguous(); + std::vector out(tensor.numel()); + auto* ptr = tensor.data_ptr(); + std::copy(ptr, ptr + tensor.numel(), out.begin()); + return out; +} + +torch::Tensor load_band_wavenumber(harp::RadiationBandOptions const& options) { + TORCH_CHECK(!options->opacities().empty(), "Band has no opacity sources"); + auto const& filename = + options->opacities().begin()->second->opacity_files().front(); + +#ifdef NETCDFOUTPUT + int fileid = harp::open_file(filename); + int varid = -1; + harp::check_nc(nc_inq_varid(fileid, "wavenumber", &varid), + "Missing required variable wavenumber"); + auto wavenumber = harp::convert_wavenumber_to_cm1( + harp::read_1d_variable(fileid, "wavenumber"), + harp::read_var_units(fileid, varid), "wavenumber"); + harp::check_nc(nc_close(fileid), "Failed to close NetCDF file"); + return wavenumber; +#else + TORCH_CHECK(false, "NetCDF support is not enabled"); +#endif +} + +std::pair, std::vector> make_bin_edges( + std::vector const& centers) { + TORCH_CHECK(centers.size() >= 2, "Need at least two wavenumber points"); + + std::vector lower(centers.size()), upper(centers.size()); + for (size_t i = 0; i < centers.size(); ++i) { + double left = 0.0; + double right = 0.0; + if (i == 0) { + left = centers[0] - 0.5 * (centers[1] - centers[0]); + } else { + left = 0.5 * (centers[i - 1] + centers[i]); + } + if (i + 1 == centers.size()) { + right = centers[i] + 0.5 * (centers[i] - centers[i - 1]); + } else { + right = 0.5 * (centers[i] + centers[i + 1]); + } + lower[i] = left; + upper[i] = right; + } + return {lower, upper}; +} + +void configure_band_grid(harp::RadiationBandOptions const& options) { + auto const wavenumber_tensor = load_band_wavenumber(options); + auto const wavenumber = tensor_to_vector(wavenumber_tensor); + auto const [wave_lower, wave_upper] = make_bin_edges(wavenumber); + std::vector weight(wavenumber.size()); + for (size_t i = 0; i < wavenumber.size(); ++i) { + weight[i] = wave_upper[i] - wave_lower[i]; + } + + options->nwave(static_cast(wavenumber.size())); + options->wavenumber(wavenumber); + options->weight(weight); + options->set_wave_lower(wave_lower); + options->set_wave_upper(wave_upper); +} + +struct ColumnState { + torch::Tensor pressure_levels; + torch::Tensor temperature_levels; + torch::Tensor pres; + torch::Tensor temp; + torch::Tensor conc; + torch::Tensor dz; +}; + +ColumnState build_column(YAML::Node const& config) { + TORCH_CHECK(config["problem"], "Missing 'problem' section"); + TORCH_CHECK(config["forcing"], "Missing 'forcing' section"); + TORCH_CHECK(config["forcing"]["const-gravity"], + "Missing 'forcing.const-gravity' section"); + + auto const problem = config["problem"]; + auto const gravity = config["forcing"]["const-gravity"]; + auto const nlyr = config["geometry"]["cells"]["nx1"].as(); + + harp::Composition composition_raw = { + {"H2", problem["xH2"].as()}, + {"He", problem["xHe"].as()}, + {"CH4", problem["xCH4"].as()}, + {"H2O", problem["xH2O"].as()}, + {"NH3", problem["xNH3"].as()}, + }; + auto const composition = harp::normalize_composition(composition_raw); + + auto const pbot = problem["Pbot"].as(); + auto const ptop = problem["Ptop"].as(); + auto const tbot = problem["Tbot"].as(); + auto const ttop = problem["Ttop"].as(); + auto const grav = std::abs(gravity["grav1"].as()); + + std::vector pressure_levels(nlyr + 1); + std::vector temperature_levels(nlyr + 1); + for (int i = 0; i <= nlyr; ++i) { + double const frac = static_cast(i) / static_cast(nlyr); + double const logp = + std::log(pbot) + frac * (std::log(ptop) - std::log(pbot)); + pressure_levels[i] = std::exp(logp); + temperature_levels[i] = tbot + frac * (ttop - tbot); + } + + std::vector pressure_layers(nlyr); + std::vector temperature_layers(nlyr); + for (int i = 0; i < nlyr; ++i) { + pressure_layers[i] = std::sqrt(pressure_levels[i] * pressure_levels[i + 1]); + temperature_layers[i] = + 0.5 * (temperature_levels[i] + temperature_levels[i + 1]); + } + + auto pressure_level_tensor = torch::tensor(pressure_levels, torch::kFloat64); + auto pres = torch::tensor(pressure_layers, torch::kFloat64).unsqueeze(0); + auto temperature_level_tensor = + torch::tensor(temperature_levels, torch::kFloat64); + auto temp = torch::tensor(temperature_layers, torch::kFloat64).unsqueeze(0); + + auto const total_molar_concentration = pres / (temp * harp::constants::Rgas); + auto conc = + torch::zeros({1, nlyr, static_cast(harp::species_names.size())}, + torch::kFloat64); + for (size_t i = 0; i < harp::species_names.size(); ++i) { + auto it = composition.find(harp::species_names[i]); + double const mixing_ratio = (it == composition.end()) ? 0.0 : it->second; + conc.index_put_({0, torch::indexing::Slice(), (int64_t)i}, + total_molar_concentration.squeeze(0) * mixing_ratio); + } + + double mean_molecular_weight = + harp::mean_molecular_weight(conc)[0].mean().item(); + TORCH_CHECK(mean_molecular_weight > 0.0, + "Mean molecular weight must be positive"); + + auto g_over_r = torch::tensor( + mean_molecular_weight * grav / harp::constants::Rgas, torch::kFloat64); + auto dz = harp::calc_dz_hypsometric(pressure_level_tensor, temp.squeeze(0), + g_over_r); + + return { + pressure_level_tensor, temperature_level_tensor, pres, temp, conc, dz}; +} + +void print_layer_summary(std::string const& label, + torch::Tensor const& attenuation, int layer_index) { + auto layer = attenuation.index({torch::indexing::Slice(), 0, layer_index, 0}); + std::cout << " " << std::setw(12) << std::left << label + << " mean[m^-1] = " << layer.mean().item() + << " max[m^-1] = " << layer.max().item() << "\n"; +} + +} // namespace + +int main(int argc, char** argv) { + std::filesystem::path exe_dir = "."; + if (argc > 0) { + exe_dir = std::filesystem::absolute(argv[0]).parent_path(); + harp::add_resource_directory(exe_dir.string(), true); + } + + auto yaml_path = (exe_dir / "jupiter_1d_rt.yaml").string(); + std::filesystem::path cli_output_dir = "."; + + for (int i = 1; i < argc; ++i) { + std::string arg = argv[i]; + if (arg == "--config") { + TORCH_CHECK(i + 1 < argc, "--config requires a file path"); + yaml_path = argv[++i]; + } else if (arg == "--output-dir") { + TORCH_CHECK(i + 1 < argc, "--output-dir requires a directory path"); + cli_output_dir = argv[++i]; + } else { + TORCH_CHECK(false, "Unknown argument: ", arg, + ". Supported options are --config and --output-dir"); + } + } + + auto yaml_dir = std::filesystem::absolute(yaml_path).parent_path(); + harp::add_resource_directory(yaml_dir.string(), true); + + auto config = YAML::LoadFile(yaml_path); + auto rad_options = harp::RadiationOptionsImpl::from_yaml(yaml_path); + TORCH_CHECK(rad_options->bands().size() == 1, + "Example expects exactly one radiation band"); + + auto band_options = rad_options->bands().front(); + configure_band_grid(band_options); + harp::RadiationBand band(band_options); + + auto column = build_column(config); + auto const reference_layer = + config["problem"]["reference_layer"].as(band_options->nlyr() / 2); + TORCH_CHECK(reference_layer >= 0 && reference_layer < band_options->nlyr(), + "Invalid reference_layer index: ", reference_layer); + + auto conc = column.conc; + auto pres = column.pres; + auto temp = column.temp; + + std::map atm; + atm["pres"] = pres; + atm["temp"] = temp; + + auto const ncol = band_options->ncol(); + auto const nwave = band_options->nwave(); + std::map bc; + bc[band_options->name() + "/albedo"] = + torch::zeros({nwave, ncol}, torch::kFloat64); + bc[band_options->name() + "/temis"] = + torch::zeros({nwave, ncol}, torch::kFloat64); + bc["btemp"] = + torch::tensor({config["problem"]["surface_temperature_k"].as( + config["problem"]["Tbot"].as())}, + torch::kFloat64); + bc["ttemp"] = + torch::tensor({config["problem"]["top_temperature_k"].as( + config["problem"]["Ttop"].as())}, + torch::kFloat64); + + auto const wavenumber = band_options->wavenumber(); + std::vector> source_attenuation; + std::vector> source_transmittance; + torch::Tensor total_attenuation; + bool first = true; + for (auto& [name, module] : band->opacities) { + auto attenuation = module.forward(conc, atm); + source_attenuation.push_back({name, attenuation}); + auto tau = (attenuation.squeeze(-1) * column.dz.view({1, 1, -1})) + .sum(-1) + .squeeze(1); + source_transmittance.push_back({name, torch::exp(-tau)}); + if (first) { + total_attenuation = attenuation.clone(); + first = false; + } else { + total_attenuation += attenuation; + } + } + + std::cout << "Band: " << band_options->name() << "\n"; + std::cout << "Layers = " << band_options->nlyr() + << ", spectral points = " << band_options->nwave() << "\n"; + std::cout << "Reference layer = " << reference_layer << "\n"; + std::cout << " pressure = " << pres[0][reference_layer].item() + << " Pa\n"; + std::cout << " temperature = " << temp[0][reference_layer].item() + << " K\n"; + + for (auto const& [name, attenuation] : source_attenuation) { + print_layer_summary(name, attenuation, reference_layer); + } + print_layer_summary("total", total_attenuation, reference_layer); + + auto band_flux = band->forward(conc, column.dz.unsqueeze(0), &bc, &atm); + auto upward = band_flux.index({0, torch::indexing::Slice(), 0}); + auto downward = band_flux.index({0, torch::indexing::Slice(), 1}); + auto net = upward - downward; + + std::cout << "\nIntegrated longwave flux profile [W/m^2]\n"; + std::cout << " surface upward = " << upward[0].item() << "\n"; + std::cout << " surface downward = " << downward[0].item() << "\n"; + std::cout << " surface net = " << net[0].item() << "\n"; + std::cout << " TOA upward = " << upward[-1].item() << "\n"; + std::cout << " TOA downward = " << downward[-1].item() << "\n"; + std::cout << " TOA net = " << net[-1].item() << "\n"; + + auto const output_dir = + std::filesystem::absolute(cli_output_dir) / + config["problem"]["output_dir"].as("jupiter_1d_rt_outputs"); + harp::write_column_profile(output_dir / "column_profile.txt", column.dz, pres, + temp, conc, band_flux); + harp::write_spectral_profile( + output_dir / "spectral_profile.txt", wavenumber, source_transmittance, + band->spectrum.index({torch::indexing::Slice(), 0, -1, 0}), + band->spectrum.index({torch::indexing::Slice(), 0, 0, 1})); + + std::cout << "\nWrote diagnostics to " << output_dir.string() << "\n"; + return 0; +} diff --git a/examples/jupiter_1d_rt.yaml b/examples/jupiter_1d_rt.yaml new file mode 100644 index 0000000..1ee65f6 --- /dev/null +++ b/examples/jupiter_1d_rt.yaml @@ -0,0 +1,67 @@ +species: + - name: H2 + composition: {H: 2} + - name: He + composition: {He: 1} + - name: CH4 + composition: {C: 1, H: 4} + - name: H2O + composition: {H: 2, O: 1} + - name: NH3 + composition: {N: 1, H: 3} + +geometry: + cells: + nx1: 100 + +opacities: + h2o_line: + type: molecule-line + data: [h2_0p9_he_0p1_ch4_0p004_h2o_0p002_nh3_0p0003_xsection_120_730K_1_100bar_20_2500cm1.nc] + species: [H2O] + + nh3_line: + type: molecule-line + data: [h2_0p9_he_0p1_ch4_0p004_h2o_0p002_nh3_0p0003_xsection_120_730K_1_100bar_20_2500cm1.nc] + species: [NH3] + + ch4_line: + type: molecule-line + data: [h2_0p9_he_0p1_ch4_0p004_h2o_0p002_nh3_0p0003_xsection_120_730K_1_100bar_20_2500cm1.nc] + species: [CH4] + + h2_h2_cia: + type: molecule-cia + data: [h2_0p9_he_0p1_ch4_0p004_h2o_0p002_nh3_0p0003_xsection_120_730K_1_100bar_20_2500cm1.nc] + species: [H2] + + h2_he_cia: + type: molecule-cia + data: [h2_0p9_he_0p1_ch4_0p004_h2o_0p002_nh3_0p0003_xsection_120_730K_1_100bar_20_2500cm1.nc] + species: [H2, He] + +bands: + - name: dump_lw + range: [20.0, 2500.0] + opacities: [h2o_line, nh3_line, ch4_line, h2_h2_cia, h2_he_cia] + solver: disort + nwave: 2481 + nstr: 4 + flags: lamber,quiet,onlyfl,planck + +forcing: + const-gravity: + grav1: -24.79 + +problem: + xH2: 0.9 + xHe: 0.1 + xCH4: 0.004 + xH2O: 0.002 + xNH3: 0.0003 + + Pbot: 100.e5 + Ptop: 1.e4 + + Tbot: 600. + Ttop: 130. diff --git a/examples/test2/amars-ck.yaml b/examples/test2/amars-ck.yaml deleted file mode 100644 index 0f4dc01..0000000 --- a/examples/test2/amars-ck.yaml +++ /dev/null @@ -1,169 +0,0 @@ -# Radiation configuration file -# -# Required the fields: -# - species: list of species -# - opacities: list of opacities -# - bands: list of bands -# - : band configuration -# -# Required fields for 'species': -# - name: species name -# - composition: species atomic composition -# -# Required fields for 'opacities': -# - type: choose between the following recognized types -# - 'rfm-lbl': RFM line-by-line table -# - 'rfm-ck': correlated-k table generated from RFM line-by-line -# - 'h2so4_simple': simple H2SO4 shortwave opacity model -# - 's8_fuller': simple S8 shortwave opacity model -# -# Optional fields for 'opacities': -# - data: list of opacity data files -# - species: list of species used by the opacity source -# -# Required fields for each band: -# - range: band wavenumber/wavelength range -# If the range is in wavenumber, the unit is cm^{-1} -# If the range is in wavelength, the unit is um -# - opacities: list of opacities used by the band -# - integration: choose between 'wavelength', 'wavenumber', or 'weight' -# If the integration is 'weight', the band is integrated -# using weights from input. Otherwise, the band is integrated -# using the trapezoidal rule along a wavelength -# or wavenumber grid (ww). -# - solver: choose between 'disort' or 'twostr' -# - flags: a string of flags for the radiative transfer solver -# separated by commas -# -# Optional fields for each band: -# - ww: a list of weights (double) for spectral bins in the band - -species: - - name: CO2 - composition: {C: 1, O: 2} - - - name: H2O - composition: {H: 2, O: 1} - - - name: SO2 - composition: {S: 1, O: 2} - - - name: H2 - composition: {H: 2} - - - name: H2SO4 - composition: {H: 2, S: 1, O: 4} - - - name: S8 - composition: {S: 8} - -opacities: - H2SO4: - type: jit - data: [grey-opacity.pt] - species: [H2SO4] - jit_kwargs: ["temp"] - nmom: 4 - - S8: - type: jit - data: [grey-opacity2.pt] - species: [S8] - jit_kwargs: ["temp"] - nmom: 4 - - CO2: - type: rfm-ck - data: [ck-h2_25-so2-3bar-.nc] - species: [CO2] - - H2O: - type: rfm-ck - data: [ck-h2_25-so2-3bar-.nc] - species: [H2O] - - SO2: - type: rfm-ck - data: [ck-h2_25-so2-3bar-.nc] - species: [SO2] - - H2: - type: rfm-ck - data: [ck-h2_25-so2-3bar-.nc] - species: [H2] - - CO2CO2CIA: - type: jit - data: [co2_co2-.pt] - species: [CO2, CO2] - - CO2H2CIA: - type: jit - data: [co2_h2-.pt] - species: [CO2, H2] - - #bands: [SW, B1, B2, B3, B4, B5, B6, B7, B8] -bands: [SW] - -SW: - range: [2000., 50000.] - opacities: [H2SO4, S8] - solver: disort - integration: wavenumber - flags: lamber,quiet,onlyfl,print-input - -B1: - range: [1., 250.] - opacities: [H2SO4, S8] - solver: disort - integration: weight - flags: lamber,quiet,onlyfl,planck - -B2: - range: [250., 438.] - opacities: [SO2, CO2, H2O, H2, CO2CO2CIA, CO2H2CIA, H2SO4, S8] - solver: disort - integration: weight - flags: lamber,quiet,onlyfl,planck - -B3: - range: [438., 675.] - opacities: [SO2, CO2, H2O, H2, CO2CO2CIA, CO2H2CIA, H2SO4, S8] - integration: weight - solver: disort - flags: lamber,quiet,onlyfl,planck - -B4: - range: [675., 1062.] - opacities: [SO2, CO2, H2O, H2, CO2CO2CIA, CO2H2CIA, H2SO4, S8] - integration: weight - solver: disort - flags: lamber,quiet,onlyfl,planck - -B5: - range: [1062., 1200.] - opacities: [SO2, CO2, H2O, H2, CO2CO2CIA, CO2H2CIA, H2SO4, S8] - integration: weight - solver: disort - flags: lamber,quiet,onlyfl,planck - -B6: - range: [1200., 1600.] - opacities: [SO2, CO2, H2O, H2, CO2CO2CIA, CO2H2CIA, H2SO4, S8] - integration: weight - solver: disort - flags: lamber,quiet,onlyfl,planck - -B7: - range: [1600., 1900.] - opacities: [SO2, CO2, H2O, H2, CO2CO2CIA, CO2H2CIA, H2SO4, S8] - integration: weight - solver: disort - flags: lamber,quiet,onlyfl,planck - -B8: - range: [1900., 2000.] - opacities: [SO2, CO2, H2O, H2, CO2CO2CIA, CO2H2CIA, H2SO4, S8] - integration: weight - solver: disort - flags: lamber,quiet,onlyfl,planck diff --git a/examples/test2/example_jit.py b/examples/test2/example_jit.py deleted file mode 100644 index 3f877f3..0000000 --- a/examples/test2/example_jit.py +++ /dev/null @@ -1,118 +0,0 @@ -import torch -import pyharp -from pyharp import ( - h2_cia_legacy, - RadiationOptions, - Radiation, - ) -import numpy as np -from pyharp.opacity import OpacityOptions, JITOpacity - -torch.set_default_dtype(torch.float64) - -class GreyOpacity(torch.nn.Module): - """ - GreyOpacity class for handling grey opacity calculations. - - The class signature and variable dimensions are strict. - The class must have a forward method that takes a concentration tensor - and returns an opacity tensor. - - The concentration vector is 3D with dimensions (ncol, nlyr, nspecies), - where ncol is the number of columns, nlyr is the number of layers, - and nspecies is the number of species. - - The opacity tensor is 4D with dimensions (nwave, ncol, nlyr, nprop), - where nwave is the number of wavelengths, ncol is the number of columns, - nlyr is the number of layers, and nprop is the number of optical properties. - - The first optical property is the total extinction cross-section [m^2/mol]. - The second optical property is the single scattering albedo. - Starting from the third, the optical properties are phase function moments - (excluding the zero-th moment). - - This class is later compiled to a TorchScript file using the `pyharp.compile` function. - """ - def __init__(self, nwave: int, nprop: int): - super().__init__() - self.nwave = nwave - self.nprop = nprop - - def forward(self, conc, temp) -> torch.Tensor: - """ - Args: - conc: Tensor of shape (ncol, nlyr, nspecies) representing concentrations. - Returns: - Tensor of shape (nwave, ncol, nlyr, nprop) representing optical properties. - """ - ncol, nlyr = conc.shape[0], conc.shape[1] - return 0.1 * torch.ones((self.nwave, ncol, nlyr, self.nprop), dtype=torch.float64) - -# compile a jit model to file -nlyr = 4 -nwave = 10 -ncol = 1 -nprop = 6 -nstr = 4 - - -model = GreyOpacity(nwave, nprop) -scripted = torch.jit.script(model) -scripted.save(f"grey-opacity.pt") - -model = GreyOpacity(nwave, nprop) -scripted = torch.jit.script(model) -scripted.save(f"grey-opacity2.pt") - -conc = torch.ones(ncol, nlyr, 6) * 0.2 - -temp = 300.0 * torch.ones(nlyr).unsqueeze(0) -print(temp) - -rad_op = RadiationOptions.from_yaml('amars-ck.yaml') - -dz = torch.ones(nlyr)*1000.0 - -for band in rad_op.bands(): - wmin = band.disort().wave_lower()[0] - wmax = band.disort().wave_upper()[0] - - band.disort().accur(1.0e-12) - band.wavenumber(list(np.linspace(wmin, wmax, nwave))) - wave = torch.tensor(band.wavenumber(), dtype=torch.float64) - print(band.wavenumber()) - bc = {} - bc[band.name() + "/fbeam"] = torch.tensor(55.).expand(nwave,ncol) - bc[band.name() + "/albedo"] = 0.3 * torch.ones((nwave,ncol)) - bc[band.name() + "/umu0"] = torch.ones((ncol,)) - -print(rad_op) -print(conc.shape) -#exit() - -atm = {} -atm['pres'] = torch.tensor([50000.,20000.,10000.,5000.]).unsqueeze(0) -atm['temp'] = temp - -rad = Radiation(rad_op) - -#just hacking these to be the right dtype -#conc = conc.float() -#dz = dz.float() -#for key in bc: -# if isinstance(bc[key], torch.Tensor): -# bc[key] = bc[key].float() -# elif isinstance(bc[key], np.ndarray): -# bc[key] = torch.from_numpy(bc[key]).float() -#for key in atm: -# if isinstance(atm[key], torch.Tensor): -# atm[key] = atm[key].float() -# elif isinstance(atm[key], np.ndarray): -# atm[key] = torch.from_numpy(atm[key]).float() - -print("dz = ", dz) -print("atm = ", atm) -print("bc = ", bc) - -netflux, downward_flux, upward_flux = rad.forward(conc, dz, bc, atm) -print(netflux) diff --git a/examples/test2/grey-opacity.pt b/examples/test2/grey-opacity.pt deleted file mode 100644 index ecd290f..0000000 Binary files a/examples/test2/grey-opacity.pt and /dev/null differ diff --git a/examples/test2/grey-opacity2.pt b/examples/test2/grey-opacity2.pt deleted file mode 100644 index 37cbcc8..0000000 Binary files a/examples/test2/grey-opacity2.pt and /dev/null differ diff --git a/python/__init__.py b/python/__init__.py index 095ee36..8484d40 100644 --- a/python/__init__.py +++ b/python/__init__.py @@ -9,7 +9,6 @@ raise from .disort import * -from .rfmlib import * from .compile import * try: diff --git a/python/csrc/pyharp.cpp b/python/csrc/pyharp.cpp index 1bfd621..3b51b25 100644 --- a/python/csrc/pyharp.cpp +++ b/python/csrc/pyharp.cpp @@ -7,6 +7,9 @@ // harp #include #include +#include +#include +#include namespace py = pybind11; @@ -54,5 +57,27 @@ PYBIND11_MODULE(pyharp, m) { return harp::deserialize_search_paths(harp::search_paths); }, py::arg("path"), py::arg("prepend") = true) - .def("find_resource", &harp::find_resource, py::arg("filename")); + .def("find_resource", &harp::find_resource, py::arg("filename")) + .def("mean_molecular_weight", &harp::mean_molecular_weight, + py::arg("conc")) + .def( + "write_column_profile", + [](const std::string& filename, torch::Tensor dz, torch::Tensor pres, + torch::Tensor temp, torch::Tensor conc, torch::Tensor band_flux) { + harp::write_column_profile(filename, dz, pres, temp, conc, + band_flux); + }, + py::arg("filename"), py::arg("dz"), py::arg("pres"), py::arg("temp"), + py::arg("conc"), py::arg("band_flux")) + .def( + "write_spectral_profile", + [](const std::string& filename, const std::vector& wavenumber, + const std::vector>& + transmittance, + torch::Tensor flux_up_toa, torch::Tensor flux_down_boa) { + harp::write_spectral_profile(filename, wavenumber, transmittance, + flux_up_toa, flux_down_boa); + }, + py::arg("filename"), py::arg("wavenumber"), py::arg("transmittance"), + py::arg("flux_up_toa"), py::arg("flux_down_boa")); } diff --git a/python/csrc/pyopacity.cpp b/python/csrc/pyopacity.cpp index cf62894..8e720bd 100644 --- a/python/csrc/pyopacity.cpp +++ b/python/csrc/pyopacity.cpp @@ -4,10 +4,11 @@ // harp #include #include +#include +#include #include #include #include -#include #include // python @@ -29,8 +30,6 @@ void bind_opacity(py::module& parent) { a->report(ss); return fmt::format("OpacityOptions(\n{})", ss.str()); }) - .def("query_wavenumber", &harp::OpacityOptionsImpl::query_wavenumber) - .def("query_weight", &harp::OpacityOptionsImpl::query_weight) .ADD_OPTION(std::string, harp::OpacityOptionsImpl, type) .ADD_OPTION(std::string, harp::OpacityOptionsImpl, bname) .ADD_OPTION(std::vector, harp::OpacityOptionsImpl, @@ -45,5 +44,7 @@ void bind_opacity(py::module& parent) { ADD_HARP_MODULE(WaveTemp, OpacityOptions, py::arg("conc"), py::arg("atm")); ADD_HARP_MODULE(MultiBand, OpacityOptions, py::arg("conc"), py::arg("atm")); ADD_HARP_MODULE(FourColumn, OpacityOptions, py::arg("conc"), py::arg("atm")); - ADD_HARP_MODULE(RFM, OpacityOptions, py::arg("conc"), py::arg("atm")); + ADD_HARP_MODULE(MoleculeLine, OpacityOptions, py::arg("conc"), + py::arg("atm")); + ADD_HARP_MODULE(MoleculeCIA, OpacityOptions, py::arg("conc"), py::arg("atm")); } diff --git a/python/csrc/pyradiation.cpp b/python/csrc/pyradiation.cpp index b4ed9d0..2986042 100644 --- a/python/csrc/pyradiation.cpp +++ b/python/csrc/pyradiation.cpp @@ -59,7 +59,11 @@ void bind_radiation(py::module& m) { .ADD_OPTION(int, harp::RadiationBandOptionsImpl, ncol) .ADD_OPTION(int, harp::RadiationBandOptionsImpl, nlyr) .ADD_OPTION(int, harp::RadiationBandOptionsImpl, nstr) - .ADD_OPTION(bool, harp::RadiationBandOptionsImpl, verbose); + .ADD_OPTION(bool, harp::RadiationBandOptionsImpl, verbose) + .def("set_wave_lower", &harp::RadiationBandOptionsImpl::set_wave_lower, + py::arg("value"), py::return_value_policy::reference) + .def("set_wave_upper", &harp::RadiationBandOptionsImpl::set_wave_upper, + py::arg("value"), py::return_value_policy::reference); auto pyRadiationOptions = py::class_( diff --git a/python/klib.py b/python/klib.py deleted file mode 100644 index fdcafe0..0000000 --- a/python/klib.py +++ /dev/null @@ -1,274 +0,0 @@ -def get_gauss_legendre(n: int): - """ - Get the Gauss-Legendre quadrature points and weights - - Parameters - ---------- - n : int - The number of quadrature points - - Returns - ------- - p : np.ndarray - The quadrature points - w : np.ndarray - The quadrature weights - """ - - p, w = np.polynomial.legendre.leggauss(n) # range is -1 to 1 - # now change it from 0 to 1 - return (p + 1.0) / 2.0, w / 2.0 - -class CorrelatedKtable: - """ - A base class to generate a correlated k-table - - Attributes - ---------- - name : str - The name of the absorbing species - """ - - def __init__(self, species: List[str]): - """ - Initialize the class - - Parameters - ---------- - name : str - The name of the absorbing species - """ - self.species = species - self.name = "-".join(species) - - def load_opacity(self, fname: str) -> None: - """ - Load the opacity data from a file. - This is a virtual function that should be implemented by the derived class. - - Parameters - ---------- - fname : str - The name of the file - """ - pass - - def make_ck_coeff( - self, - kcoeff: np.ndarray, - ilayer: int, - nbins: int, - npoints: int, - log_opacity: bool = True, - ) -> np.ndarray: - """ - Make the correlated k-table axis and sort the k-coefficients - - Parameters - ---------- - kcoeff : np.ndarray - The absorption coefficients with shape (nwave, nlayer) - ilayer : int - The layer index to use for the bin_divides - nbins : int - The number of bins to divide the spectral band - npoints : int - The number of points in each bin - log_opacity : bool - If True, the absorption coefficients are in log scale - - Returns - ------- - ckcoeff : np.ndarray - The correlated k-coefficients with shape (nbins * npoints, nlayer) - """ - bin_divides = np.zeros(nbins + 1) - bin_divides[0] = 0.0 - bin_divides[-1] = 1.0 - - # use ilayer to determine the bin_divides - if log_opacity: - lnkmax = kcoeff[:, ilayer].max() - lnkmin = kcoeff[:, ilayer].min() - else: - lnkmax = np.log(kcoeff[:, ilayer].max()) - lnkmin = np.log(kcoeff[:, ilayer].min()) - - nwave = kcoeff.shape[0] - nlayer = kcoeff.shape[1] - - for i in range(nlayer): - kcoeff[:, i] = np.sort(kcoeff[:, i]) - - for i in range(1, nbins): - lnk = lnkmin + (lnkmax - lnkmin) * i / nbins - if log_opacity: - bin_divides[i] = np.searchsorted(kcoeff[:, ilayer], lnk) / nwave - else: - bin_divides[i] = np.searchsorted(kcoeff[:, ilayer], exp(lnk)) / nwave - - # print('bin_divides:', bin_divides) - - gaxis = np.zeros(nbins * npoints) - weights = np.zeros(nbins * npoints) - ckcoeff = np.zeros((nbins * npoints, nlayer)) - - gg, ww = get_gauss_legendre(npoints) - for i in range(nbins): - gaxis[i * npoints : (i + 1) * npoints] = ( - gg * (bin_divides[i + 1] - bin_divides[i]) + bin_divides[i] - ) - weights[i * npoints : (i + 1) * npoints] = ww * ( - bin_divides[i + 1] - bin_divides[i] - ) - - for j in range(nlayer): - kcoeff_func = interp1d(np.arange(nwave), kcoeff[:, j]) - ckcoeff[:, j] = kcoeff_func(gaxis * (nwave - 1)) - - self.gaxis = gaxis - self.weights = weights - - return ckcoeff - - def write_opacity(self, fname: str): - """ - Write the correlated k-table to a file - - Parameters - ---------- - fname : str - The name of the file - """ - ncfile = Dataset(fname, "w") - ncfile.createDimension("gaxis", len(self.gaxis)) - dim = ncfile.createVariable("gaxis", "f8", ("gaxis",)) - dim[:] = self.gaxis - dim.long_name = "gaussian quadrature points" - dim.units = "1" - - ncfile.createDimension("Wavenumber", len(self.gaxis)) - dim = ncfile.createVariable("Wavenumber", "f8", ("Wavenumber",)) - dim[:] = self.wave - dim.long_name = "gaussian quadrature equivalent wavenumber" - dim.units = "1/cm" - - ncfile.createDimension("weights", len(self.weights)) - dim = ncfile.createVariable("weights", "f8", ("weights",)) - dim[:] = self.weights - dim.long_name = "gaussian quadrature weights" - dim.units = "1" - - ncfile.createDimension("Pressure", len(self.pres)) - dim = ncfile.createVariable("Pressure", "f8", ("Pressure",)) - dim[:] = self.pres - dim.long_name = "reference pressure" - dim.units = "pa" - - ncfile.createDimension("TempGrid", len(self.temp_grid)) - dim = ncfile.createVariable("TempGrid", "f8", ("TempGrid",)) - dim[:] = self.temp_grid - dim.long_name = "temperature anomaly grid" - dim.units = "K" - - var = ncfile.createVariable("Temperature", "f8", ("Pressure")) - var[:] = self.temp - var.long_name = "reference temperature" - var.units = "K" - - for name in self.species: - var = ncfile.createVariable( - name, "f8", ("Wavenumber", "Pressure", "TempGrid") - ) - var[:] = self.ckcoeff[name] - var.long_name = "correlated k-coefficients" - var.units = self.kunits[name] - - ncfile.close() - print("Correlated k-table written to", fname) - - -class HitranCorrelatedKtable(CorrelatedKtable): - """ - Derived class to generate a correlated k-table from HITRAN line-by-line opacity - """ - - def load_opacity(self, fname: str): - """ - Load the opacity data from a file. Overrides the base class method. - - Parameters - ---------- - fname : str - The name of the file - """ - data = Dataset(fname, "r") - self.kcoeff = {} - self.kunits = {} - - for name in self.species: - self.kcoeff[name] = data.variables[name][:] - self.kunits[name] = data.variables[name].units - self.pres = data.variables["Pressure"][:] - self.temp = data.variables["Temperature"][:] - self.temp_grid = data.variables["TempGrid"][:] - - def make_cktable(self, wmin: float, wmax: float, nbins: int = 1, npoints: int = 50): - """ - Make the correlated k-table. This function will call make_ck_axis for each temperature - grid point. - - Parameters - ---------- - nbins : int - The number of bins to divide the spectral band - npoints : int - The number of points in each bin - """ - nlayer = len(self.pres) - ntemp = len(self.temp_grid) - - self.ckcoeff = {} - for name in self.species: - self.ckcoeff[name] = np.zeros((nbins * npoints, nlayer, ntemp)) - for i in range(ntemp): - self.ckcoeff[name][:, :, i] = self.make_ck_coeff( - self.kcoeff[name][:, :, i], nlayer // 2, nbins, npoints - ) - self.wave = wmin + self.gaxis * (wmax - wmin) - - -def run_cktable_one_band(bname: str): - band = radiation_band(bname, config) - - species = list(map(str, config[bname]["opacity"])) - wmin, wmax = band.get_range() - ab_ck = HitranCorrelatedKtable(species) - ab_ck.load_opacity(opacity_input + "-" + bname + ".nc") - ab_ck.make_cktable(wmin, wmax) - ab_ck.write_opacity(opacity_output + "-" + bname + ".nc") - -def run_ktable_one_band(bname: str): - band = radiation_band(bname, config) - - nspec = band.get_num_specgrids() - wmin, wmax = band.get_range() - wres = (wmax - wmin) / (nspec - 1) - - # atmospheric properties - num_layers = len(atm["HGT"]) - wav_grid = (wmin, wmax, wres) - - species = list(map(str, config[bname]["opacity"])) - driver = create_rfm_driver(wav_grid, tem_grid, species, hitran_file) - - # write rfm atmosphere file to file - write_rfm_atm(atm, rundir=bname) - write_rfm_drv(driver, rundir=bname) - - # run rfm and write kcoeff file - pwd = os.getcwd() - run_rfm(rundir=bname) - write_ktable( - opacity_output + "-" + bname, species, atm, wav_grid, tem_grid, basedir=pwd - ) diff --git a/python/opacity.pyi b/python/opacity.pyi index e46a9ae..de5b8da 100644 --- a/python/opacity.pyi +++ b/python/opacity.pyi @@ -17,7 +17,7 @@ class OpacityOptions: Examples: >>> import torch >>> from pyharp.opacity import OpacityOptions - >>> op = OpacityOptions().type('rfm-lbl') + >>> op = OpacityOptions().type('molecule-line') """ def __init__(self) -> None: @@ -59,7 +59,7 @@ class OpacityOptions: """ Set the type of the opacity source format. - Valid options are: ``jit``, ``rfm-lbl``, ``rfm-ck``, ``fourcolumn``, ``wavetemp``, ``multiband-ck``, ``helios``. + Valid options are: ``jit``, ``molecule-line``, ``molecule-cia``, ``fourcolumn``, ``wavetemp``, ``multiband-ck``, ``helios``. Args: value (str): type of the opacity source @@ -70,7 +70,7 @@ class OpacityOptions: Examples: >>> import torch >>> from pyharp.opacity import OpacityOptions - >>> op = OpacityOptions().type('rfm-lbl') + >>> op = OpacityOptions().type('molecule-line') >>> print(op) """ ... @@ -538,14 +538,14 @@ class FourColumn: """ ... -class RFM: +class MoleculeLine: """ - Line-by-line absorption data computed by RFM. + Molecular line absorption read from a NetCDF dump. Examples: >>> import torch - >>> from pyharp.opacity import RFM, OpacityOptions - >>> op = RFM(OpacityOptions()) + >>> from pyharp.opacity import MoleculeLine, OpacityOptions + >>> op = MoleculeLine(OpacityOptions()) """ options: OpacityOptions @@ -558,7 +558,7 @@ class RFM: @overload def __init__(self, options: OpacityOptions) -> None: """ - Create a RFM instance. + Create a MoleculeLine instance. Args: options (OpacityOptions): Opacity options @@ -593,7 +593,7 @@ class RFM: def forward(self, conc: torch.Tensor, atm: dict[str, torch.Tensor]) -> torch.Tensor: """ - Calculate opacity using RFM line-by-line absorption data. + Calculate attenuation using NetCDF line cross sections. Args: conc (torch.Tensor): concentration of the species in mol/m^3 @@ -609,7 +609,81 @@ class RFM: where nwave is the number of wavelengths, ncol is the number of columns, nlyr is the number of layers. - The last dimension is the optical properties arranged - in the order of attenuation [1/m], single scattering albedo and scattering phase function. + The last dimension is the attenuation coefficient [1/m]. + """ + ... + +class MoleculeCIA: + """ + Collision-induced absorption read from a NetCDF dump. + + Examples: + >>> import torch + >>> from pyharp.opacity import MoleculeCIA, OpacityOptions + >>> op = MoleculeCIA(OpacityOptions()) + """ + + options: OpacityOptions + + @overload + def __init__(self) -> None: + """Construct a new default module.""" + ... + + @overload + def __init__(self, options: OpacityOptions) -> None: + """ + Create a MoleculeCIA instance. + + Args: + options (OpacityOptions): Opacity options + """ + ... + + def __repr__(self) -> str: ... + + def module(self, name: str): + """ + Get a submodule by name. + + Args: + name (str): name of the submodule + + Returns: + The submodule + """ + ... + + def buffer(self, name: str) -> torch.Tensor: + """ + Get a buffer by name. + + Args: + name (str): name of the buffer + + Returns: + torch.Tensor: the buffer tensor + """ + ... + + def forward(self, conc: torch.Tensor, atm: dict[str, torch.Tensor]) -> torch.Tensor: + """ + Calculate attenuation using NetCDF CIA binary coefficients. + + Args: + conc (torch.Tensor): concentration of the species in mol/m^3 + atm (dict[str, torch.Tensor]): atmospheric parameters + + Either 'wavelength' or 'wavenumber' must be provided + if 'wavelength' is provided, the unit is um. + if 'wavenumber' is provided, the unit is cm^{-1}. + + Returns: + torch.Tensor: + The shape of the output tensor is (nwave, ncol, nlyr, 1), + where nwave is the number of wavelengths, + ncol is the number of columns, + nlyr is the number of layers. + The last dimension is the attenuation coefficient [1/m]. """ ... diff --git a/python/pyharp.pyi b/python/pyharp.pyi index 6d45aa5..5f10bb5 100644 --- a/python/pyharp.pyi +++ b/python/pyharp.pyi @@ -107,6 +107,60 @@ def find_resource(filename: str) -> str: """ ... +def mean_molecular_weight(conc: torch.Tensor) -> torch.Tensor: + """ + Calculate mean molecular weight from species molar concentrations. + + Args: + conc (torch.Tensor): concentration tensor with species on the last axis + and units of mol/m^3 + + Returns: + torch.Tensor: mean molecular weight [kg/mol] with the species axis reduced + """ + ... + +def write_column_profile( + filename: str, + dz: torch.Tensor, + pres: torch.Tensor, + temp: torch.Tensor, + conc: torch.Tensor, + band_flux: torch.Tensor, +) -> None: + """ + Write a formatted 1D column profile text table. + + Args: + filename (str): output path + dz (torch.Tensor): layer thickness [m], shape (nlyr,) + pres (torch.Tensor): pressure [Pa], shape (ncol, nlyr) + temp (torch.Tensor): temperature [K], shape (ncol, nlyr) + conc (torch.Tensor): molar concentration [mol/m^3], shape (ncol, nlyr, nspecies) + band_flux (torch.Tensor): band flux [W/m^2], shape (ncol, nlyr+1, 2) + """ + ... + +def write_spectral_profile( + filename: str, + wavenumber: list[float], + transmittance: list[tuple[str, torch.Tensor]], + flux_up_toa: torch.Tensor, + flux_down_boa: torch.Tensor, +) -> None: + """ + Write a formatted spectral profile text table. + + Args: + filename (str): output path + wavenumber (list[float]): wavenumber grid [cm^-1] + transmittance (list[tuple[str, torch.Tensor]]): named total transmittance + arrays, each tensor shape (nwave,) + flux_up_toa (torch.Tensor): upward spectral flux at TOA [W/m^2/cm^-1], shape (nwave,) + flux_down_boa (torch.Tensor): downward spectral flux at BOA [W/m^2/cm^-1], shape (nwave,) + """ + ... + # Radiation functions @overload def bbflux_wavenumber(wave: torch.Tensor, temp: float, ncol: int = 1) -> torch.Tensor: @@ -432,6 +486,30 @@ class RadiationBandOptions: """ ... + def set_wave_lower(self, value: list[float]) -> "RadiationBandOptions": + """ + Set solver lower bin edges for the active band solver. + + Args: + value (list[float]): lower bin edges [cm^-1] + + Returns: + RadiationBandOptions: class object + """ + ... + + def set_wave_upper(self, value: list[float]) -> "RadiationBandOptions": + """ + Set solver upper bin edges for the active band solver. + + Args: + value (list[float]): upper bin edges [cm^-1] + + Returns: + RadiationBandOptions: class object + """ + ... + @overload def verbose(self) -> bool: """Get verbose flag.""" diff --git a/python/rfmlib.py b/python/rfmlib.py deleted file mode 100644 index dab4471..0000000 --- a/python/rfmlib.py +++ /dev/null @@ -1,297 +0,0 @@ -# Purpose: Library for RFM calculations - -import torch -import os, subprocess, shutil -import numpy as np -from collections import OrderedDict -from typing import List, Tuple, Dict - -def create_rfm_driver( - wav_grid: Tuple[float, float, float], - tem_grid: Tuple[int, float, float], - absorbers: List[str], - hitran_file: str, -) -> Dict[str, str]: - """ - Create a RFM driver file. - - Parameters - ---------- - wav_grid : Tuple[float, float, float] - Wavenumber grid by minimum, maximum and resolution. - tem_grid : Tuple[int, float, float] - Temperature grid by number of points, minimum and maximum. - absorbers : List - A list of absorbers. - hitran_file : str - Path to HITRAN file. - - Returns - ------- - driver : Dict[str, str] - A dictionary containing the driver file content. - """ - driver = OrderedDict( - [ - ("*HDR", "Header for rfm"), - ("*FLG", "TAB CTM"), - ("*SPC", "%.4f %.4f %.4f" % wav_grid), - ("*GAS", " ".join(absorbers)), - ("*ATM", "rfm.atm"), - ("*DIM", "PLV \n %d %.4f %.4f" % tem_grid), - ("*TAB", "tab_*.txt"), - ("*HIT", hitran_file), - ("*END", ""), - ] - ) - return driver - -def write_rfm_atm(atm: Dict[str, np.ndarray], rundir: str=".") -> None: - """ - Write RFM atmosphere to file. - - Parameters - ---------- - atm : Dict[str, np.ndarray] - A dictionary containing the atmosphere - rundir : str - Directory to write the file. Default is current directory. - - Returns - ------- - None - """ - print(f"# Creating {rundir}/rfm.atm ...") - num_layers = atm["HGT"].shape[0] - if not os.path.exists(f"{rundir}"): - os.makedirs(f"{rundir}") - - with open(f"{rundir}/rfm.atm", "w") as file: - file.write("%d\n" % num_layers) - file.write("*HGT [km]\n") - for i in range(num_layers): # m -> km - file.write("%.8g " % (atm["HGT"][i] / 1.0e3,)) - file.write("\n*PRE [mb]\n") - for i in range(num_layers): # pa -> mb - file.write("%.8g " % (atm["PRE"][i] / 100.0,)) - file.write("\n*TEM [K]\n") - for i in range(num_layers): - file.write("%.8g " % atm["TEM"][i]) - for name, val in atm.items(): - if name in ["IDX", "HGT", "PRE", "TEM"]: - continue - file.write("\n*" + name + " [ppmv]\n") - for j in range(num_layers): # mol/mol -> ppmv - file.write("%.8g " % (val[j] * 1.0e6,)) - file.write("\n*END") - print(f"# {rundir}/rfm.atm written.") - -def write_rfm_drv(driver: Dict[str, str], rundir: str=".") -> None: - """ - Write RFM driver to file. - - Parameters - ---------- - driver : Dict[str, str] - A dictionary containing the driver file content. - rundir : str - Directory to write the file. Default is current directory - - Returns - ------- - None - """ - print(f"# Creating {rundir}/rfm.drv ...") - if not os.path.exists(f"{rundir}"): - os.makedirs(f"{rundir}") - - with open(f"{rundir}/rfm.drv", "w") as file: - for sec in driver: - if driver[sec] != None: - file.write(sec + "\n") - file.write(" " * 4 + driver[sec] + "\n") - print(f"# {rundir}/rfm.drv written.") - -def run_rfm(rundir: str=".") -> None: - """ - Call to run RFM. - - Parameters - ---------- - None - - Returns - ------- - None - """ - pwd = os.getcwd() - if not os.path.exists(rundir): - os.makedirs(rundir) - - if rundir != ".": - os.chdir(os.path.join(pwd, rundir)) - - with open(f"rfm.runlog", "w") as file: - process = subprocess.Popen( - [f"{pwd}/rfm.release"], stdout=file, stderr=subprocess.STDOUT - ) - process.communicate() - - #for line in iter(process.stdout.readline, b""): - # decode the byte string and end='' to avoid double newlines - # print(line.decode(), end="") - - -def create_netcdf_input( - fname: str, - absorbers: List[str], - atm: Dict[str, np.ndarray], - wmin: float, - wmax: float, - wres: float, - tnum: int, - tmin: float, - tmax: float, -) -> str: - """ - Create an input file for writing kcoeff table to netCDF format - - Parameters - ---------- - fname : str - Name of the file. - absorbers : list - A list of absorbers. - atm : Dict[str, np.ndarray] - A dictionary containing the atmosphere. - wmin : float - Minimum wavenumber. - wmax : float - Maximum wavenumber. - wres : float - Wavenumber resolution. - tnum : int - Number of temperature points. - tmin : float - Minimum temperature. - tmax : float - Maximum temperature. - - Returns - ------- - fname : str - Name of the input file for netCDf - """ - print(f"# Creating {fname}.in ...") - - with open(f"{fname}.in", "w") as file: - file.write("# Molecular absorber\n") - file.write("%d\n" % len(absorbers)) - file.write(" ".join(absorbers) + "\n") - file.write("# Molecule data files\n") - for ab in absorbers: - file.write("%-40s\n" % (f"tab_" + ab.lower() + ".txt",)) - file.write("# Wavenumber range\n") - file.write( - "%-14.6g%-14.6g%-14.6g\n" % (wmin, wmax, int((wmax - wmin) / wres) + 1) - ) - file.write("# Relative temperature range\n") - file.write("%-14.6g%-14.6g%-14.6g\n" % (tmin, tmax, tnum)) - file.write("# Number of vertical levels\n") - file.write("%d\n" % len(atm["TEM"])) - file.write("# Temperature\n") - for i in range(len(atm["TEM"])): - file.write("%-14.6g" % atm["TEM"][-(i + 1)]) - if (i + 1) % 10 == 0: - file.write("\n") - if (i + 1) % 10 != 0: - file.write("\n") - file.write("# Pressure\n") - for i in range(len(atm["PRE"])): - file.write("%-14.6g" % atm["PRE"][-(i + 1)]) - if (i + 1) % 10 == 0: - file.write("\n") - if (i + 1) % 10 != 0: - file.write("\n") - print(f"# {fname}.in written.") - return f"{fname}.in" - -def write_ktable( - fname: str, - absorbers: List[str], - atm: Dict[str, np.ndarray], - wav_grid: Tuple[float, float, float], - tem_grid: Tuple[int, float, float], - basedir: str=".", -) -> None: - """ - Write kcoeff table to netCDF file. - - Parameters - ---------- - fname : str - Name of the file. - absorbers : List - A list of absorbers. - atm : Dict[str, np.ndarray] - A dictionary containing the atmosphere. - wav_grid : Tuple[float, float, float] - Wavenumber grid by minimum, maximum and resolution. - tem_grid : Tuple[int, float, float] - Temperature grid by number of points, minimum and maximum. - - Returns - ------- - None - """ - inpfile = create_netcdf_input(fname, absorbers, atm, *wav_grid, *tem_grid) - - process = subprocess.Popen( - [f"{basedir}/kcoeff.release", "-i", inpfile, "-o", f"{fname}.nc"], - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - ) - - for line in iter(process.stdout.readline, b""): - # decode the byte string and end='' to avoid double newlines - print(line.decode(), end="") - - process.communicate() - - pwd = os.getcwd() - shutil.move(f"{fname}.nc", f"{basedir}/{fname}.nc") - print(f"# {fname}.nc written.") - -def read_rfm_atm(filename): - data = {} - with open(filename, "r") as f: - lines = f.readlines() - - current_key = None - current_values = [] - - for line in lines: - line = line.strip() - if not line: - continue - - if line.startswith("*"): - if current_key is not None: - # Save the previous section - value = torch.tensor(current_values, dtype=torch.float32) - data[current_key] = value - current_values = [] - - if line == "*END": - break - - current_key = line.split()[0][1:] # remove '*' - else: - current_values.extend(map(float, line.split())) - - # Save the last variable block if not already saved - if current_key is not None and current_key not in data: - value = torch.tensor(current_values, dtype=torch.float32) - data[current_key] = value - - return data diff --git a/python/run_cktable.py b/python/run_cktable.py deleted file mode 100644 index fc91d5e..0000000 --- a/python/run_cktable.py +++ /dev/null @@ -1,28 +0,0 @@ -#! /usr/bin/env python3 -import sys - -sys.path.append("../python") -sys.path.append(".") - -import canoe -from canoe import load_configure -from canoe.harp import radiation_band -from netCDF4 import Dataset -from scipy.interpolate import interp1d -from multiprocessing import Pool, cpu_count -from typing import List -import numpy as np - - -if __name__ == "__main__": - canoe.start() - - opacity_input = "amarsw-lbl" - opacity_output = "amarsw-ck" - opacity_config = "amarsw-op.yaml" - max_threads = cpu_count() - - config = load_configure(opacity_config) - pool = Pool(max_threads) - bnames = list(map(str, config["bands"])) - pool.map(run_cktable_one_band, bnames) diff --git a/python/run_ktable_amars.py b/python/run_ktable_amars.py deleted file mode 100644 index 1763335..0000000 --- a/python/run_ktable_amars.py +++ /dev/null @@ -1,33 +0,0 @@ -#! /usr/bin/env python3 -import sys - -sys.path.append("../python") -sys.path.append(".") - -import canoe -from canoe import load_configure, find_resource -from canoe.harp import radiation_band -from atm_profile_utils import read_atm_profile -from multiprocessing import Pool, cpu_count -from typing import Tuple -from rfmlib import * - - - - -if __name__ == "__main__": - canoe.start() - - atm_profile = "amarsw.atm" - opacity_config = "amarsw-op.yaml" - opacity_output = "amarsw-lbl" - tem_grid = (5, -50, 50) - max_threads = cpu_count() - - hitran_file = find_resource("HITRAN2020.par") - atm = read_atm_profile(atm_profile) - config = load_configure(opacity_config) - - pool = Pool(max_threads) - bnames = list(map(str, config["bands"])) - pool.map(run_ktable_one_band, bnames) diff --git a/python/spectra/atm_overview_cli.py b/python/spectra/atm_overview.py similarity index 90% rename from python/spectra/atm_overview_cli.py rename to python/spectra/atm_overview.py index 2f005c7..845f4dd 100644 --- a/python/spectra/atm_overview_cli.py +++ b/python/spectra/atm_overview.py @@ -4,11 +4,14 @@ import argparse from collections.abc import Iterator -from concurrent.futures import ProcessPoolExecutor +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor from dataclasses import dataclass import json import os from pathlib import Path +import pickle +import subprocess +import sys import tempfile os.environ.setdefault("MPLCONFIGDIR", str(Path(tempfile.gettempdir()) / "spectra_matplotlib")) @@ -27,11 +30,21 @@ supported_hitran_cia_pairs, supported_hitran_species_names, ) -from .hitran_cia import load_cia_dataset -from .hitran_lines import build_line_provider, download_hitran_lines, load_hitran_line_list -from .molecule_plot_cli import _add_legend_if_needed, _apply_positive_log_scale, _mask_nonpositive, _plot_transmittance_panel, _style_shared_axis +from .hitran_cia_utils import load_cia_dataset +from .hitran_molecule_plot import ( + _add_legend_if_needed, + _apply_positive_log_scale, + _mask_nonpositive, + _plot_transmittance_panel, + _style_shared_axis, +) +from .hitran_molecule_utils import ( + build_line_provider, + download_hitran_lines, + load_hitran_line_list, +) from .mt_ckd_h2o import compute_mt_ckd_h2o_continuum_cross_section -from .shared_cli import process_pool_context +from .utils import build_band_from_range, default_hitran_dir, parse_wn_range, process_pool_context from .spectrum import AbsorptionSpectrum, number_density_cm3_from_pressure_temperature from .transmittance import compute_transmittance_spectrum @@ -75,7 +88,7 @@ def _canonical_mixture_species_names() -> dict[str, str]: return mapping -def _parse_composition(value: str) -> dict[str, float]: +def parse_composition(value: str) -> dict[str, float]: canonical_names = _canonical_mixture_species_names() entries: list[tuple[str, float]] = [] for chunk in str(value).split(","): @@ -108,29 +121,6 @@ def _parse_composition(value: str) -> dict[str, float]: raise ValueError("composition fractions must sum to a positive value") return {species_name: value / total_fraction for species_name, value in totals.items()} - -def _parse_wn_range(value: str) -> tuple[float, float]: - lower_text, sep, upper_text = str(value).partition(",") - if not sep: - raise argparse.ArgumentTypeError("wn-range must have the form min,max") - try: - lower = float(lower_text) - upper = float(upper_text) - except ValueError as exc: - raise argparse.ArgumentTypeError("wn-range must contain numeric min,max values") from exc - if upper < lower: - raise argparse.ArgumentTypeError("wn-range max must be >= min") - return lower, upper - - -def _build_band(*, wn_min: float, wn_max: float, resolution: float) -> SpectralBandConfig: - if resolution <= 0.0: - raise ValueError("resolution must be positive") - if wn_max < wn_min: - raise ValueError("wn-range max must be >= min") - return SpectralBandConfig("single_state", float(wn_min), float(wn_max), float(resolution)) - - def _build_species_config( *, species_name: str, @@ -168,11 +158,11 @@ def _line_supported_species(composition: dict[str, float]) -> tuple[tuple[str, f def compute_mixture_overview_products(args: argparse.Namespace, *, wn_range: tuple[float, float]) -> MixtureOverviewProducts: if args.path_length_km <= 0.0: raise ValueError("path-length-km must be positive") - composition = _parse_composition(args.composition) + composition = parse_composition(args.composition) broadening_composition = parse_broadening_composition( getattr(args, "broadening_composition", None) or composition ) - band = _build_band(wn_min=wn_range[0], wn_max=wn_range[1], resolution=float(args.resolution)) + band = build_band_from_range(wn_range, float(args.resolution)) grid = band.grid() temperature_k = float(args.temperature_k) pressure_pa = float(args.pressure_bar) * 1.0e5 @@ -534,13 +524,13 @@ def _state_pairs(args: argparse.Namespace) -> list[tuple[float, float]]: def build_atm_overview_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="Generate a 4-row atmospheric opacity overview PDF for a gas mixture.") - parser.add_argument("--hitran-dir", type=Path, default=Path("hitran")) + parser.add_argument("--hitran-dir", type=Path, default=default_hitran_dir()) parser.add_argument("--composition", required=True) parser.add_argument("--temperature-k", type=float, default=300.0) parser.add_argument("--pressure-bar", type=float, default=1.0) parser.add_argument("--path-length-km", type=float, default=1.0) parser.add_argument("--resolution", type=float, default=1.0) - parser.add_argument("--wn-range", dest="wn_ranges", action="append", type=_parse_wn_range, required=True) + parser.add_argument("--wn-range", dest="wn_ranges", action="append", type=parse_wn_range, required=True) parser.add_argument("--broadening-composition", default=None) parser.add_argument("--refresh-hitran", action="store_true") parser.add_argument("--refresh-cia", action="store_true") @@ -587,7 +577,7 @@ def run_atm_overview(args: argparse.Namespace) -> None: manifest = { "composition_input": str(args.composition), - "composition_normalized": _parse_composition(args.composition), + "composition_normalized": parse_composition(args.composition), "temperature_k": _selected_temperatures(args), "pressure_bar": _selected_pressure_bars(args), "path_length_km": float(args.path_length_km), @@ -616,5 +606,45 @@ def _parallel_mixture_overview_products( return max_workers = min(len(tasks), os.cpu_count() or 1) ctx = process_pool_context() - with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as executor: - yield from executor.map(_compute_mixture_overview_product_task, tasks) + try: + with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as executor: + yield from executor.map(_compute_mixture_overview_product_task, tasks) + except PermissionError: + yield from _parallel_mixture_overview_products_via_subprocess(tasks, max_workers=max_workers) + + +def _subprocess_worker_entry(task_pickle_path: str, result_pickle_path: str, worker_name: str) -> None: + with Path(task_pickle_path).open("rb") as handle: + task = pickle.load(handle) + worker = globals()[worker_name] + result = worker(task) + with Path(result_pickle_path).open("wb") as handle: + pickle.dump(result, handle, protocol=pickle.HIGHEST_PROTOCOL) + + +def _run_task_via_subprocess(task: tuple[argparse.Namespace, tuple[float, float]]) -> MixtureOverviewProducts: + code = ( + "from pyharp.spectra.atm_overview import _subprocess_worker_entry; " + "import sys; " + "_subprocess_worker_entry(sys.argv[1], sys.argv[2], sys.argv[3])" + ) + with tempfile.TemporaryDirectory(prefix="pyharp_overview_task_") as tmpdir: + task_path = Path(tmpdir) / "task.pkl" + result_path = Path(tmpdir) / "result.pkl" + with task_path.open("wb") as handle: + pickle.dump(task, handle, protocol=pickle.HIGHEST_PROTOCOL) + subprocess.run( + [sys.executable, "-c", code, str(task_path), str(result_path), "_compute_mixture_overview_product_task"], + check=True, + ) + with result_path.open("rb") as handle: + return pickle.load(handle) + + +def _parallel_mixture_overview_products_via_subprocess( + tasks: list[tuple[argparse.Namespace, tuple[float, float]]], + *, + max_workers: int, +) -> Iterator[MixtureOverviewProducts]: + with ThreadPoolExecutor(max_workers=max_workers) as executor: + yield from executor.map(_run_task_via_subprocess, tasks) diff --git a/python/spectra/config.py b/python/spectra/config.py index 11a1b2b..3fb781b 100644 --- a/python/spectra/config.py +++ b/python/spectra/config.py @@ -19,7 +19,7 @@ class SpectralBandConfig: resolution_cm1: float = 1.0 def grid(self) -> np.ndarray: - """Return the lower-inclusive, upper-exclusive wavenumber grid.""" + """Return the wavenumber grid with both endpoints included.""" if self.resolution_cm1 <= 0: raise ValueError("resolution_cm1 must be positive") span_cm1 = self.wavenumber_max_cm1 - self.wavenumber_min_cm1 @@ -34,9 +34,9 @@ def grid(self) -> np.ndarray: ): raise ValueError( "wavenumber range must be an integer multiple of resolution_cm1 " - "to produce a lower-inclusive, upper-exclusive grid" + "to produce an endpoint-inclusive grid" ) - return self.wavenumber_min_cm1 + np.arange(count, dtype=np.float64) * self.resolution_cm1 + return self.wavenumber_min_cm1 + np.arange(count + 1, dtype=np.float64) * self.resolution_cm1 @dataclass(frozen=True) diff --git a/python/spectra/dataset_io.py b/python/spectra/dataset_io.py index f461ead..d0e2732 100644 --- a/python/spectra/dataset_io.py +++ b/python/spectra/dataset_io.py @@ -2,9 +2,9 @@ from __future__ import annotations +from contextlib import contextmanager +import os from pathlib import Path -import shutil -import tempfile import numpy as np import xarray as xr @@ -12,6 +12,8 @@ DEFAULT_NETCDF_ENGINE = "scipy" WAVENUMBER_ATTRS = {"long_name": "wavenumber", "units": "cm^-1"} +HDF5_USE_FILE_LOCKING_ENV = "HDF5_USE_FILE_LOCKING" +HDF5_USE_FILE_LOCKING_DISABLED = "FALSE" def clean_var_token(value: object) -> str: @@ -86,12 +88,20 @@ def combine_band_datasets(datasets: list[xr.Dataset], *, wn_ranges: list[tuple[f return combined -def write_dataset_via_tmp(dataset: xr.Dataset, output_path: Path, *, engine: str = DEFAULT_NETCDF_ENGINE) -> None: - output_path.parent.mkdir(parents=True, exist_ok=True) - fd, tmp_name = tempfile.mkstemp(prefix="pyharp_", suffix=".nc", dir="/tmp") - Path(tmp_name).unlink(missing_ok=True) +@contextmanager +def hdf5_file_locking_disabled(): + previous_value = os.environ.get(HDF5_USE_FILE_LOCKING_ENV) + os.environ[HDF5_USE_FILE_LOCKING_ENV] = HDF5_USE_FILE_LOCKING_DISABLED try: - dataset.to_netcdf(tmp_name, engine=engine) - shutil.move(tmp_name, output_path) + yield finally: - Path(tmp_name).unlink(missing_ok=True) + if previous_value is None: + os.environ.pop(HDF5_USE_FILE_LOCKING_ENV, None) + else: + os.environ[HDF5_USE_FILE_LOCKING_ENV] = previous_value + + +def write_dataset(dataset: xr.Dataset, output_path: Path, *, engine: str = DEFAULT_NETCDF_ENGINE) -> None: + output_path.parent.mkdir(parents=True, exist_ok=True) + with hdf5_file_locking_disabled(): + dataset.to_netcdf(output_path, engine=engine) diff --git a/python/spectra/dump_cli.py b/python/spectra/dump_cli.py index 5eab090..e80f373 100644 --- a/python/spectra/dump_cli.py +++ b/python/spectra/dump_cli.py @@ -3,16 +3,19 @@ from __future__ import annotations import argparse -from concurrent.futures import ProcessPoolExecutor +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor import os from pathlib import Path +import pickle +import subprocess import sys +import tempfile from textwrap import dedent import numpy as np import xarray as xr -from .atm_overview_cli import _parse_composition, compute_mixture_overview_products +from .atm_overview import compute_mixture_overview_products, parse_composition from .config import SpectroscopyConfig, cia_database_for_model, parse_broadening_composition, resolve_hitran_cia_filename, resolve_hitran_cia_pair from .dataset_io import ( DEFAULT_NETCDF_ENGINE, @@ -20,14 +23,15 @@ add_wavenumber_attrs, build_state_attrs, clean_var_token, - write_dataset_via_tmp, + write_dataset, ) -from .hitran_cia import load_cia_dataset -from .hitran_lines import build_line_provider, download_hitran_lines +from .hitran_cia_utils import load_cia_dataset +from .hitran_molecule_utils import build_line_provider, download_hitran_lines from .orton_xiz_cia import load_orton_xiz_cia_dataset, resolve_orton_xiz_cia_filename -from .output_names import _clean_token, _format_value -from .shared_cli import ( +from .utils import ( HelpFormatter, + _clean_token, + _format_value, build_band, default_hitran_dir, default_orton_xiz_cia_dir, @@ -580,7 +584,7 @@ def _write_xsection_dataset( output_path.parent.mkdir(parents=True, exist_ok=True) dataset = _xsection_dataset(spectrum, species_name=species_name, secondary_component=secondary_component, wn_range=wn_range) try: - write_dataset_via_tmp(dataset, output_path, engine=DEFAULT_NETCDF_ENGINE) + write_dataset(dataset, output_path, engine=DEFAULT_NETCDF_ENGINE) finally: dataset.close() @@ -797,7 +801,7 @@ def _compute_transmission_band(task: tuple[str, argparse.Namespace]) -> tuple[xr def _composition_species_names(composition: str) -> tuple[str, ...]: - return tuple(_parse_composition(composition).keys()) + return tuple(parse_composition(composition).keys()) def _parallel_band_results( @@ -809,8 +813,52 @@ def _parallel_band_results( return [worker(task) for task in tasks] max_workers = min(len(tasks), os.cpu_count() or 1) ctx = process_pool_context() - with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as executor: - return list(executor.map(worker, tasks)) + try: + with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as executor: + return list(executor.map(worker, tasks)) + except PermissionError: + return _parallel_band_results_via_subprocess(tasks, worker=worker, max_workers=max_workers) + + +def _subprocess_worker_entry(task_pickle_path: str, result_pickle_path: str, worker_name: str) -> None: + with Path(task_pickle_path).open("rb") as handle: + task = pickle.load(handle) + worker = globals()[worker_name] + result = worker(task) + with Path(result_pickle_path).open("wb") as handle: + pickle.dump(result, handle, protocol=pickle.HIGHEST_PROTOCOL) + + +def _run_task_via_subprocess(task: tuple[str, argparse.Namespace], *, worker_name: str) -> tuple[xr.Dataset, str | None]: + code = ( + "from pyharp.spectra.dump_cli import _subprocess_worker_entry; " + "import sys; " + "_subprocess_worker_entry(sys.argv[1], sys.argv[2], sys.argv[3])" + ) + with tempfile.TemporaryDirectory(prefix="pyharp_dump_task_") as tmpdir: + task_path = Path(tmpdir) / "task.pkl" + result_path = Path(tmpdir) / "result.pkl" + with task_path.open("wb") as handle: + pickle.dump(task, handle, protocol=pickle.HIGHEST_PROTOCOL) + subprocess.run( + [sys.executable, "-c", code, str(task_path), str(result_path), worker_name], + check=True, + ) + with result_path.open("rb") as handle: + return pickle.load(handle) + + +def _parallel_band_results_via_subprocess( + tasks: list[tuple[str, argparse.Namespace]], + *, + worker, + max_workers: int, +) -> list[tuple[xr.Dataset, str | None]]: + worker_name = getattr(worker, "__name__", None) + if worker_name is None or globals().get(worker_name) is not worker: + raise ValueError("worker must be a module-level dump_cli function for subprocess fallback") + with ThreadPoolExecutor(max_workers=max_workers) as executor: + return list(executor.map(lambda task: _run_task_via_subprocess(task, worker_name=worker_name), tasks)) def _stack_state_grid_datasets( @@ -974,7 +1022,7 @@ def main() -> None: ): output_path = _output_path_for_wn_range(args, wn_range=wn_range, suffix=".nc") try: - write_dataset_via_tmp(dataset, output_path, engine=DEFAULT_NETCDF_ENGINE) + write_dataset(dataset, output_path, engine=DEFAULT_NETCDF_ENGINE) finally: dataset.close() print(f"Wrote NetCDF: {output_path}") @@ -992,7 +1040,7 @@ def main() -> None: ): output_path = _output_path_for_wn_range(args, wn_range=wn_range, suffix=".nc") try: - write_dataset_via_tmp(dataset, output_path, engine=DEFAULT_NETCDF_ENGINE) + write_dataset(dataset, output_path, engine=DEFAULT_NETCDF_ENGINE) finally: dataset.close() print(f"Wrote NetCDF: {output_path}") diff --git a/python/spectra/cia_plot_cli.py b/python/spectra/hitran_cia_plot.py similarity index 52% rename from python/spectra/cia_plot_cli.py rename to python/spectra/hitran_cia_plot.py index 365563e..147a804 100644 --- a/python/spectra/cia_plot_cli.py +++ b/python/spectra/hitran_cia_plot.py @@ -1,64 +1,153 @@ -"""Shared CLI helpers for CIA plotting scripts.""" +"""Plotting and CLI helpers for HITRAN collision-induced absorption.""" from __future__ import annotations import argparse from pathlib import Path + import numpy as np -from .config import SpectralBandConfig, cia_database_for_model, resolve_hitran_cia_filename, resolve_hitran_cia_pair -from .hitran_cia import ( +from .blackbody import compute_normalized_blackbody_curve +from .config import cia_database_for_model, resolve_hitran_cia_filename +from .hitran_cia_utils import ( + CiaDataset, + compute_cia_attenuation_m1, + compute_cia_transmission, load_cia_dataset, - plot_cia_attenuation_coefficient, - plot_cia_cross_section, - plot_cia_transmission, ) from .orton_xiz_cia import load_orton_xiz_cia_dataset, resolve_orton_xiz_cia_filename -from .shared_cli import default_orton_xiz_cia_dir - - -def _parse_wn_range(value: str) -> tuple[float, float]: - lower_text, sep, upper_text = str(value).partition(",") - if not sep: - raise argparse.ArgumentTypeError("wn-range must have the form min,max") - try: - lower = float(lower_text) - upper = float(upper_text) - except ValueError as exc: - raise argparse.ArgumentTypeError("wn-range must contain numeric min,max values") from exc - if upper < lower: - raise argparse.ArgumentTypeError("wn-range max must be >= min") - return lower, upper - +from .utils import build_grid, default_hitran_dir, default_orton_xiz_cia_dir, parse_wn_range + + +def plot_cia_cross_section( + dataset: CiaDataset, + *, + temperature_k: float, + wavenumber_grid_cm1: np.ndarray, + figure_path: Path, +) -> np.ndarray: + """Plot a CIA binary absorption coefficient spectrum in cm^5 molecule^-2.""" + import matplotlib.pyplot as plt + + wavenumber_grid_cm1 = np.asarray(wavenumber_grid_cm1, dtype=np.float64) + binary_xsec = dataset.interpolate_to_grid(temperature_k, wavenumber_grid_cm1) + positive = binary_xsec[binary_xsec > 0.0] + if positive.size == 0: + raise ValueError("No positive CIA coefficients were found on the requested grid.") + + figure_path.parent.mkdir(parents=True, exist_ok=True) + fig, ax = plt.subplots(figsize=(11, 6)) + ax.plot(wavenumber_grid_cm1, binary_xsec, color="tab:blue", linewidth=1.25) + ax.set_yscale("log") + ax.set_xlabel("Wavenumber [cm$^{-1}$]") + ax.set_ylabel("Binary absorption coefficient [cm$^5$ / molecule$^2$]") + ax.set_title(f"{dataset.pair} CIA at {temperature_k:.1f} K") + ax.grid(True, which="both", alpha=0.3) + fig.tight_layout() + fig.savefig(figure_path, dpi=150) + plt.close(fig) + return binary_xsec + + +def plot_cia_attenuation_coefficient( + dataset: CiaDataset, + *, + temperature_k: float, + pressure_pa: float, + wavenumber_grid_cm1: np.ndarray, + figure_path: Path, +) -> np.ndarray: + """Plot a CIA attenuation coefficient spectrum in 1/m.""" + import matplotlib.pyplot as plt + + wavenumber_grid_cm1 = np.asarray(wavenumber_grid_cm1, dtype=np.float64) + attenuation_m1 = compute_cia_attenuation_m1( + dataset, + temperature_k=temperature_k, + pressure_pa=pressure_pa, + wavenumber_grid_cm1=wavenumber_grid_cm1, + ) + positive = attenuation_m1[attenuation_m1 > 0.0] + if positive.size == 0: + raise ValueError("No positive CIA attenuation coefficients were found on the requested grid.") + + figure_path.parent.mkdir(parents=True, exist_ok=True) + fig, ax = plt.subplots(figsize=(11, 6)) + ax.plot(wavenumber_grid_cm1, attenuation_m1, color="tab:orange", linewidth=1.25) + ax.set_yscale("log") + ax.set_xlabel("Wavenumber [cm$^{-1}$]") + ax.set_ylabel("Attenuation coefficient [1/m]") + ax.set_title(f"{dataset.pair} CIA attenuation at {temperature_k:.1f} K and {pressure_pa / 1.0e5:.3f} bar") + ax.grid(True, which="both", alpha=0.3) + fig.tight_layout() + fig.savefig(figure_path, dpi=150) + plt.close(fig) + return attenuation_m1 + + +def plot_cia_transmission( + dataset: CiaDataset, + *, + temperature_k: float, + pressure_pa: float, + path_length_m: float, + wavenumber_grid_cm1: np.ndarray, + figure_path: Path, +) -> np.ndarray: + """Plot CIA transmission over a fixed path length.""" + import matplotlib.pyplot as plt + + wavenumber_grid_cm1 = np.asarray(wavenumber_grid_cm1, dtype=np.float64) + transmission = compute_cia_transmission( + dataset, + temperature_k=temperature_k, + pressure_pa=pressure_pa, + path_length_m=path_length_m, + wavenumber_grid_cm1=wavenumber_grid_cm1, + ) -def _wn_bounds(args: argparse.Namespace) -> tuple[float, float]: - return tuple(args.wn_range) + figure_path.parent.mkdir(parents=True, exist_ok=True) + fig, ax = plt.subplots(figsize=(11, 6)) + ax.plot(wavenumber_grid_cm1, transmission, color="tab:green", linewidth=1.25) + ax.plot( + wavenumber_grid_cm1, + compute_normalized_blackbody_curve( + wavenumber_cm1=wavenumber_grid_cm1, + temperature_k=temperature_k, + ), + color="black", + linestyle="--", + linewidth=1.1, + label="Blackbody", + ) + ax.set_xlabel("Wavenumber [cm$^{-1}$]") + ax.set_ylabel("Transmission") + ax.set_ylim(0.0, 1.01) + ax.set_title( + f"{dataset.pair} CIA transmission at {temperature_k:.1f} K, " + f"{pressure_pa / 1.0e5:.3f} bar, L={path_length_m / 1000.0:.3f} km" + ) + ax.grid(True, which="both", alpha=0.3) + ax.legend() + fig.tight_layout() + fig.savefig(figure_path, dpi=150) + plt.close(fig) + return transmission def _add_common_arguments(parser: argparse.ArgumentParser) -> None: - parser.add_argument("--hitran-dir", type=Path, default=Path("hitran")) + parser.add_argument("--hitran-dir", type=Path, default=default_hitran_dir()) parser.add_argument("--cia-dir", type=Path, default=None) parser.add_argument("--cia-model", choices=("auto", "2011", "2018", "xiz", "orton"), default="auto") parser.add_argument("--h2-state", choices=("eq", "nm"), default="eq") parser.add_argument("--filename", default=None) parser.add_argument("--pair", default="H2-H2") parser.add_argument("--temperature-k", type=float, default=300.0) - parser.add_argument("--wn-range", type=_parse_wn_range, default=(20.0, 10000.0)) + parser.add_argument("--wn-range", type=parse_wn_range, default=(20.0, 10000.0)) parser.add_argument("--resolution", type=float, default=1.0) parser.add_argument("--refresh", action="store_true") -def _validate_and_build_grid(args: argparse.Namespace) -> np.ndarray: - wn_min, wn_max = _wn_bounds(args) - band = SpectralBandConfig( - name="cia_plot", - wavenumber_min_cm1=wn_min, - wavenumber_max_cm1=wn_max, - resolution_cm1=args.resolution, - ) - return band.grid() - - def _resolve_filename(args: argparse.Namespace) -> str: if args.filename: return str(args.filename) @@ -106,7 +195,7 @@ def main_binary() -> None: def run_binary(args: argparse.Namespace) -> None: - grid = _validate_and_build_grid(args) + grid = build_grid(args, name="cia_plot") dataset = _load_dataset(args) values = plot_cia_cross_section( dataset, @@ -132,7 +221,7 @@ def main_attenuation() -> None: def run_attenuation(args: argparse.Namespace) -> None: - grid = _validate_and_build_grid(args) + grid = build_grid(args, name="cia_plot") dataset = _load_dataset(args) values = plot_cia_attenuation_coefficient( dataset, @@ -162,7 +251,7 @@ def main_transmission() -> None: def run_transmission(args: argparse.Namespace) -> None: if args.path_length_km <= 0.0: raise ValueError("path-length-km must be positive") - grid = _validate_and_build_grid(args) + grid = build_grid(args, name="cia_plot") dataset = _load_dataset(args) values = plot_cia_transmission( dataset, diff --git a/python/spectra/hitran_cia.py b/python/spectra/hitran_cia_utils.py similarity index 69% rename from python/spectra/hitran_cia.py rename to python/spectra/hitran_cia_utils.py index 94d9762..2307d19 100644 --- a/python/spectra/hitran_cia.py +++ b/python/spectra/hitran_cia_utils.py @@ -1,4 +1,4 @@ -"""Download and parse HITRAN collision-induced absorption files.""" +"""Utility helpers for HITRAN collision-induced absorption data.""" from __future__ import annotations @@ -10,7 +10,6 @@ import numpy as np -from .blackbody import compute_normalized_blackbody_curve from .config import SpectroscopyConfig K_BOLTZMANN = 1.380649e-23 @@ -233,36 +232,6 @@ def load_cia_dataset( return parse_cia_file(path, pair) -def plot_cia_cross_section( - dataset: CiaDataset, - *, - temperature_k: float, - wavenumber_grid_cm1: np.ndarray, - figure_path: Path, -) -> np.ndarray: - """Plot a CIA binary absorption coefficient spectrum in cm^5 molecule^-2.""" - import matplotlib.pyplot as plt - - wavenumber_grid_cm1 = np.asarray(wavenumber_grid_cm1, dtype=np.float64) - binary_xsec = dataset.interpolate_to_grid(temperature_k, wavenumber_grid_cm1) - positive = binary_xsec[binary_xsec > 0.0] - if positive.size == 0: - raise ValueError("No positive CIA coefficients were found on the requested grid.") - - figure_path.parent.mkdir(parents=True, exist_ok=True) - fig, ax = plt.subplots(figsize=(11, 6)) - ax.plot(wavenumber_grid_cm1, binary_xsec, color="tab:blue", linewidth=1.25) - ax.set_yscale("log") - ax.set_xlabel("Wavenumber [cm$^{-1}$]") - ax.set_ylabel("Binary absorption coefficient [cm$^5$ / molecule$^2$]") - ax.set_title(f"{dataset.pair} CIA at {temperature_k:.1f} K") - ax.grid(True, which="both", alpha=0.3) - fig.tight_layout() - fig.savefig(figure_path, dpi=150) - plt.close(fig) - return binary_xsec - - def compute_cia_attenuation_m1( dataset: CiaDataset, *, @@ -277,42 +246,6 @@ def compute_cia_attenuation_m1( return np.asarray(binary_xsec * number_density_cm3**2 * 100.0, dtype=np.float64) -def plot_cia_attenuation_coefficient( - dataset: CiaDataset, - *, - temperature_k: float, - pressure_pa: float, - wavenumber_grid_cm1: np.ndarray, - figure_path: Path, -) -> np.ndarray: - """Plot a CIA attenuation coefficient spectrum in 1/m.""" - import matplotlib.pyplot as plt - - wavenumber_grid_cm1 = np.asarray(wavenumber_grid_cm1, dtype=np.float64) - attenuation_m1 = compute_cia_attenuation_m1( - dataset, - temperature_k=temperature_k, - pressure_pa=pressure_pa, - wavenumber_grid_cm1=wavenumber_grid_cm1, - ) - positive = attenuation_m1[attenuation_m1 > 0.0] - if positive.size == 0: - raise ValueError("No positive CIA attenuation coefficients were found on the requested grid.") - - figure_path.parent.mkdir(parents=True, exist_ok=True) - fig, ax = plt.subplots(figsize=(11, 6)) - ax.plot(wavenumber_grid_cm1, attenuation_m1, color="tab:orange", linewidth=1.25) - ax.set_yscale("log") - ax.set_xlabel("Wavenumber [cm$^{-1}$]") - ax.set_ylabel("Attenuation coefficient [1/m]") - ax.set_title(f"{dataset.pair} CIA attenuation at {temperature_k:.1f} K and {pressure_pa / 1.0e5:.3f} bar") - ax.grid(True, which="both", alpha=0.3) - fig.tight_layout() - fig.savefig(figure_path, dpi=150) - plt.close(fig) - return attenuation_m1 - - def compute_cia_transmission( dataset: CiaDataset, *, @@ -331,53 +264,3 @@ def compute_cia_transmission( wavenumber_grid_cm1=wavenumber_grid_cm1, ) return np.exp(-attenuation_m1 * float(path_length_m)) - - -def plot_cia_transmission( - dataset: CiaDataset, - *, - temperature_k: float, - pressure_pa: float, - path_length_m: float, - wavenumber_grid_cm1: np.ndarray, - figure_path: Path, -) -> np.ndarray: - """Plot CIA transmission over a fixed path length.""" - import matplotlib.pyplot as plt - - wavenumber_grid_cm1 = np.asarray(wavenumber_grid_cm1, dtype=np.float64) - transmission = compute_cia_transmission( - dataset, - temperature_k=temperature_k, - pressure_pa=pressure_pa, - path_length_m=path_length_m, - wavenumber_grid_cm1=wavenumber_grid_cm1, - ) - - figure_path.parent.mkdir(parents=True, exist_ok=True) - fig, ax = plt.subplots(figsize=(11, 6)) - ax.plot(wavenumber_grid_cm1, transmission, color="tab:green", linewidth=1.25) - ax.plot( - wavenumber_grid_cm1, - compute_normalized_blackbody_curve( - wavenumber_cm1=wavenumber_grid_cm1, - temperature_k=temperature_k, - ), - color="black", - linestyle="--", - linewidth=1.1, - label="Blackbody", - ) - ax.set_xlabel("Wavenumber [cm$^{-1}$]") - ax.set_ylabel("Transmission") - ax.set_ylim(0.0, 1.01) - ax.set_title( - f"{dataset.pair} CIA transmission at {temperature_k:.1f} K, " - f"{pressure_pa / 1.0e5:.3f} bar, L={path_length_m / 1000.0:.3f} km" - ) - ax.grid(True, which="both", alpha=0.3) - ax.legend() - fig.tight_layout() - fig.savefig(figure_path, dpi=150) - plt.close(fig) - return transmission diff --git a/python/spectra/molecule_plot_cli.py b/python/spectra/hitran_molecule_plot.py similarity index 67% rename from python/spectra/molecule_plot_cli.py rename to python/spectra/hitran_molecule_plot.py index ffdbf92..3d99033 100644 --- a/python/spectra/molecule_plot_cli.py +++ b/python/spectra/hitran_molecule_plot.py @@ -1,56 +1,81 @@ -"""Shared CLI helpers for molecular absorption plotting scripts.""" +"""Plotting and CLI helpers for HITRAN molecular line workflows.""" from __future__ import annotations import argparse from collections.abc import Iterator -from concurrent.futures import ProcessPoolExecutor +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor import os from pathlib import Path +import pickle +import subprocess +import sys import tempfile -os.environ.setdefault("MPLCONFIGDIR", str(Path(tempfile.gettempdir()) / "spectra_matplotlib")) -import matplotlib - -matplotlib.use("Agg") import matplotlib.pyplot as plt from matplotlib.backends.backend_pdf import PdfPages import numpy as np from .blackbody import compute_normalized_blackbody_curve -from .config import SpectroscopyConfig, SpectralBandConfig, parse_broadening_composition, resolve_hitran_cia_pair -from .hitran_cia import load_cia_dataset -from .hitran_lines import LineDatabase, build_line_provider, download_hitran_lines, load_hitran_line_list, plot_hitran_line_positions -from .shared_cli import process_pool_context -from .spectrum import _resolve_continuum_sources, compute_absorption_spectrum_from_sources, plot_absorption_spectrum, plot_attenuation_spectrum +from .config import parse_broadening_composition +from .hitran_molecule_utils import ( + _build_band_and_config, + compute_requested_absorption_spectrum, + compute_requested_transmittance, + load_requested_cia_dataset, + selected_pressure_bars, + selected_temperatures, + state_pairs, + LineDatabase, + build_line_provider, + download_hitran_lines, + load_hitran_line_list, +) +from .utils import default_hitran_dir, parse_wn_range, process_pool_context +from .spectrum import plot_absorption_spectrum, plot_attenuation_spectrum from .transmittance import compute_transmittance_spectrum, plot_transmittance_spectrum -def _parse_wn_range(value: str) -> tuple[float, float]: - lower_text, sep, upper_text = str(value).partition(",") - if not sep: - raise argparse.ArgumentTypeError("wn-range must have the form min,max") - try: - lower = float(lower_text) - upper = float(upper_text) - except ValueError as exc: - raise argparse.ArgumentTypeError("wn-range must contain numeric min,max values") from exc - if upper < lower: - raise argparse.ArgumentTypeError("wn-range max must be >= min") - return lower, upper - +def plot_hitran_line_positions( + line_list, + figure_path: Path, + *, + wavenumber_min_cm1: float | None = None, + wavenumber_max_cm1: float | None = None, + min_line_strength: float = 1.0e-27, +) -> None: + """Plot discrete HITRAN line positions and line intensities on log-log axes.""" + figure_path.parent.mkdir(parents=True, exist_ok=True) + threshold = float(min_line_strength) + positive = line_list.line_intensity[line_list.line_intensity >= threshold] + if positive.size == 0: + raise ValueError("Line list does not contain any intensities above the minimum threshold.") + ymin = threshold -def _wn_bounds(args: argparse.Namespace) -> tuple[float, float]: - return tuple(args.wn_range) + fig, ax = plt.subplots(figsize=(12, 6)) + ax.vlines(line_list.wavenumber_cm1, ymin, line_list.line_intensity, color="0.2", alpha=0.2, linewidth=0.35) + ax.scatter(line_list.wavenumber_cm1, line_list.line_intensity, s=4.0, color="tab:blue", alpha=0.7, linewidths=0.0) + ax.set_xscale("log") + ax.set_yscale("log") + ax.set_ylim(ymin, max(float(np.max(positive)) * 1.2, ymin * 10.0)) + if wavenumber_min_cm1 is not None and wavenumber_max_cm1 is not None: + ax.set_xlim(float(wavenumber_min_cm1), float(wavenumber_max_cm1)) + ax.set_xlabel("Wavenumber [cm$^{-1}$]") + ax.set_ylabel("Line intensity $S$ [cm$^{-1}$/ (molecule cm$^{-2}$)]") + ax.set_title(f"{line_list.species_name} HITRAN line positions and intensities") + ax.grid(True, which="both", alpha=0.25) + fig.tight_layout() + fig.savefig(figure_path, dpi=150) + plt.close(fig) def _add_common_arguments(parser: argparse.ArgumentParser, *, include_state: bool = True) -> None: - parser.add_argument("--hitran-dir", type=Path, default=Path("hitran")) + parser.add_argument("--hitran-dir", type=Path, default=default_hitran_dir()) parser.add_argument("--species", default="H2O") if include_state: parser.add_argument("--temperature-k", type=float, default=300.0) parser.add_argument("--pressure-bar", type=float, default=1.0) - parser.add_argument("--wn-range", type=_parse_wn_range, default=(20.0, 2500.0)) + parser.add_argument("--wn-range", type=parse_wn_range, default=(20.0, 2500.0)) parser.add_argument("--resolution", type=float, default=1.0) parser.add_argument("--refresh-hitran", action="store_true") parser.add_argument("--broadening-composition", default=None) @@ -62,23 +87,6 @@ def _add_external_cia_arguments(parser: argparse.ArgumentParser, *, required: bo parser.add_argument("--refresh-cia", action="store_true") -def _build_band_and_config(args: argparse.Namespace) -> tuple[SpectralBandConfig, SpectroscopyConfig]: - wn_min, wn_max = _wn_bounds(args) - if args.resolution <= 0.0: - raise ValueError("resolution must be positive") - if wn_max < wn_min: - raise ValueError("wn-range max must be >= min") - band = SpectralBandConfig("single_state", wn_min, wn_max, args.resolution) - config = SpectroscopyConfig( - output_path=Path("output") / "unused.nc", - hitran_cache_dir=args.hitran_dir, - species_name=args.species, - broadening_composition=parse_broadening_composition(getattr(args, "broadening_composition", None)), - refresh_hitran=args.refresh_hitran, - ) - return band, config - - def _print_positive_summary(values: np.ndarray, label: str, unit: str) -> None: positive = np.asarray(values, dtype=np.float64) positive = positive[positive > 0.0] @@ -93,69 +101,6 @@ def _print_bounded_summary(values: np.ndarray, label: str) -> None: print(f"{label} min/max: {bounded.min():.3e} .. {bounded.max():.3e}") -def _resolve_cia_pair(args: argparse.Namespace, config: SpectroscopyConfig) -> str: - return str(args.cia_pair or f"{config.hitran_species.name}-{config.hitran_species.name}") - - -def _resolve_self_component_filename(args: argparse.Namespace, config: SpectroscopyConfig) -> str | None: - explicit = getattr(args, "cia_filename", None) - if explicit: - return str(explicit) - explicit_pair = getattr(args, "cia_pair", None) - if explicit_pair: - return resolve_hitran_cia_pair(explicit_pair).filename - return config.hitran_species.cia_filename - - -def _load_requested_cia_dataset(args: argparse.Namespace, config: SpectroscopyConfig): - cia_filename = _resolve_self_component_filename(args, config) - if cia_filename is None: - return None - return load_cia_dataset( - cache_dir=args.hitran_dir, - filename=cia_filename, - pair=_resolve_cia_pair(args, config), - refresh=bool(args.refresh_cia), - ) - - -_MISSING = object() - - -def _compute_requested_absorption_spectrum( - args: argparse.Namespace, - *, - line_db: LineDatabase | None = None, - cia_dataset=_MISSING, -): - band, config = _build_band_and_config(args) - temperature_k = float(args.temperature_k) - pressure_pa = float(args.pressure_bar) * 1.0e5 - if cia_dataset is _MISSING: - cia_dataset = _load_requested_cia_dataset(args, config) - line_db = line_db or download_hitran_lines(config, band) - line_provider = build_line_provider(config, line_db) - grid = band.grid() - cia_cross_section_cm2_molecule = None - if cia_dataset is None: - cia_dataset, cia_cross_section_cm2_molecule = _resolve_continuum_sources( - config=config, - wavenumber_grid_cm1=grid, - temperature_k=temperature_k, - pressure_pa=pressure_pa, - ) - spectrum = compute_absorption_spectrum_from_sources( - species_name=config.hitran_species.name, - wavenumber_grid_cm1=grid, - temperature_k=temperature_k, - pressure_pa=pressure_pa, - line_provider=line_provider, - cia_dataset=cia_dataset, - cia_cross_section_cm2_molecule=cia_cross_section_cm2_molecule, - ) - return band, config, spectrum, line_provider - - def _print_spectrum_summary(spectrum, *, include_components: bool = False) -> None: print(f"Species: {spectrum.species_name}") print(f"Grid points: {spectrum.wavenumber_cm1.size}") @@ -188,7 +133,7 @@ def main_xsection() -> None: def run_xsection(args: argparse.Namespace) -> None: - _, _, spectrum, line_provider = _compute_requested_absorption_spectrum(args) + _, _, spectrum, line_provider = compute_requested_absorption_spectrum(args) plot_absorption_spectrum(spectrum, args.figure) print(f"Broadening: {line_provider.broadening_summary()}") print(f"Species: {spectrum.species_name}") @@ -199,16 +144,10 @@ def run_xsection(args: argparse.Namespace) -> None: def build_attenuation_parser(*, require_external_cia: bool = False, description: str | None = None, default_figure: Path | None = None) -> argparse.ArgumentParser: - parser = argparse.ArgumentParser( - description=description or "Compute and plot molecular attenuation coefficient." - ) + parser = argparse.ArgumentParser(description=description or "Compute and plot molecular attenuation coefficient.") _add_common_arguments(parser) _add_external_cia_arguments(parser, required=require_external_cia) - parser.add_argument( - "--figure", - type=Path, - default=default_figure or Path("output/molecule_attenuation_300K_1bar.png"), - ) + parser.add_argument("--figure", type=Path, default=default_figure or Path("output/molecule_attenuation_300K_1bar.png")) return parser @@ -218,7 +157,7 @@ def main_attenuation() -> None: def run_attenuation(args: argparse.Namespace) -> None: - _, _, spectrum, line_provider = _compute_requested_absorption_spectrum(args) + _, _, spectrum, line_provider = compute_requested_absorption_spectrum(args) plot_attenuation_spectrum(spectrum, args.figure) print(f"Broadening: {line_provider.broadening_summary()}") _print_spectrum_summary(spectrum, include_components=True) @@ -234,38 +173,21 @@ def main_attenuation_with_cia() -> None: def build_transmission_parser(*, require_external_cia: bool = False, description: str | None = None, default_figure: Path | None = None) -> argparse.ArgumentParser: - parser = argparse.ArgumentParser( - description=description or "Compute and plot molecular transmission over a fixed path length." - ) + parser = argparse.ArgumentParser(description=description or "Compute and plot molecular transmission over a fixed path length.") _add_common_arguments(parser) _add_external_cia_arguments(parser, required=require_external_cia) parser.add_argument("--path-length-km", type=float, default=1.0) - parser.add_argument( - "--figure", - type=Path, - default=default_figure or Path("output/molecule_transmission_300K_1bar_1km.png"), - ) + parser.add_argument("--figure", type=Path, default=default_figure or Path("output/molecule_transmission_300K_1bar_1km.png")) return parser -def _compute_requested_transmittance(args: argparse.Namespace): - if args.path_length_km <= 0.0: - raise ValueError("path-length-km must be positive") - _, _, spectrum, line_provider = _compute_requested_absorption_spectrum(args) - transmittance = compute_transmittance_spectrum( - spectrum=spectrum, - path_length_m=float(args.path_length_km) * 1000.0, - ) - return spectrum, transmittance, line_provider - - def main_transmission() -> None: args = build_transmission_parser().parse_args() run_transmission(args) def run_transmission(args: argparse.Namespace) -> None: - _, transmittance, line_provider = _compute_requested_transmittance(args) + _, transmittance, line_provider = compute_requested_transmittance(args) plot_transmittance_spectrum(transmittance, args.figure) print(f"Broadening: {line_provider.broadening_summary()}") _print_transmittance_summary(transmittance, include_components=True) @@ -358,13 +280,7 @@ def _make_overview_header_text(*, spectrum, transmittance, band) -> str: ) -def _plot_line_positions_panel( - ax, - line_list, - *, - xlim: tuple[float, float], - min_line_strength: float = 1.0e-27, -) -> None: +def _plot_line_positions_panel(ax, line_list, *, xlim: tuple[float, float], min_line_strength: float = 1.0e-27) -> None: mask = ( (line_list.wavenumber_cm1 >= xlim[0]) & (line_list.wavenumber_cm1 <= xlim[1]) @@ -389,18 +305,8 @@ def _plot_line_positions_panel( def _plot_cross_section_panel(ax, spectrum, *, xlim: tuple[float, float], show_components: bool, component_label: str) -> None: if show_components: ax.plot(spectrum.wavenumber_cm1, _mask_nonpositive(spectrum.sigma_line_cm2_molecule), label="Line", linewidth=1.0) - ax.plot( - spectrum.wavenumber_cm1, - _mask_nonpositive(spectrum.sigma_cia_cm2_molecule), - label=component_label, - linewidth=1.0, - ) - ax.plot( - spectrum.wavenumber_cm1, - _mask_nonpositive(spectrum.sigma_total_cm2_molecule), - label="Total" if show_components else "Cross section", - linewidth=1.4, - ) + ax.plot(spectrum.wavenumber_cm1, _mask_nonpositive(spectrum.sigma_cia_cm2_molecule), label=component_label, linewidth=1.0) + ax.plot(spectrum.wavenumber_cm1, _mask_nonpositive(spectrum.sigma_total_cm2_molecule), label="Total" if show_components else "Cross section", linewidth=1.4) if not _apply_positive_log_scale( ax, [ @@ -436,18 +342,8 @@ def _plot_secondary_component_panel(ax, *, cia_dataset, spectrum, temperature_k: def _plot_attenuation_panel(ax, spectrum, *, xlim: tuple[float, float], show_components: bool, component_label: str) -> None: if show_components: ax.plot(spectrum.wavenumber_cm1, _mask_nonpositive(spectrum.attenuation_line_m1), label="Line", linewidth=1.0) - ax.plot( - spectrum.wavenumber_cm1, - _mask_nonpositive(spectrum.attenuation_cia_m1), - label=component_label, - linewidth=1.0, - ) - ax.plot( - spectrum.wavenumber_cm1, - _mask_nonpositive(spectrum.attenuation_total_m1), - label="Total" if show_components else "Attenuation", - linewidth=1.4, - ) + ax.plot(spectrum.wavenumber_cm1, _mask_nonpositive(spectrum.attenuation_cia_m1), label=component_label, linewidth=1.0) + ax.plot(spectrum.wavenumber_cm1, _mask_nonpositive(spectrum.attenuation_total_m1), label="Total" if show_components else "Attenuation", linewidth=1.4) if not _apply_positive_log_scale( ax, [ @@ -466,18 +362,8 @@ def _plot_attenuation_panel(ax, spectrum, *, xlim: tuple[float, float], show_com def _plot_transmittance_panel(ax, transmittance, *, xlim: tuple[float, float], show_components: bool, component_label: str) -> None: if show_components: ax.plot(transmittance.wavenumber_cm1, transmittance.transmittance_line, label="Line", linewidth=1.0) - ax.plot( - transmittance.wavenumber_cm1, - transmittance.transmittance_cia, - label=component_label, - linewidth=1.0, - ) - ax.plot( - transmittance.wavenumber_cm1, - transmittance.transmittance_total, - label="Total" if show_components else "Transmittance", - linewidth=1.4, - ) + ax.plot(transmittance.wavenumber_cm1, transmittance.transmittance_cia, label=component_label, linewidth=1.0) + ax.plot(transmittance.wavenumber_cm1, transmittance.transmittance_total, label="Total" if show_components else "Transmittance", linewidth=1.4) ax.plot( transmittance.wavenumber_cm1, compute_normalized_blackbody_curve( @@ -498,13 +384,7 @@ def _plot_transmittance_panel(ax, transmittance, *, xlim: tuple[float, float], s def _render_overview_page(fig, axes_flat, *, band, config, spectrum, transmittance, line_list, cia_dataset) -> None: xlim = (band.wavenumber_min_cm1, band.wavenumber_max_cm1) - fig.subplots_adjust( - left=1.0 / 8.5, - right=1.0 - 1.0 / 8.5, - bottom=1.0 / 11.0, - top=1.0 - 1.0 / 11.0, - hspace=0.5, - ) + fig.subplots_adjust(left=1.0 / 8.5, right=1.0 - 1.0 / 8.5, bottom=1.0 / 11.0, top=1.0 - 1.0 / 11.0, hspace=0.5) fig.text( 1.0 / 8.5, 1.0 - 0.45 / 11.0, @@ -517,21 +397,10 @@ def _render_overview_page(fig, axes_flat, *, band, config, spectrum, transmittan has_secondary_component = bool(np.any(np.asarray(spectrum.sigma_cia_cm2_molecule, dtype=np.float64) > 0.0)) component_label = _resolve_overview_component_label(cia_dataset=cia_dataset, spectrum=spectrum) show_components = has_secondary_component - _plot_line_positions_panel( - axes_flat[0], - line_list, - xlim=xlim, - min_line_strength=config.min_line_strength, - ) + _plot_line_positions_panel(axes_flat[0], line_list, xlim=xlim, min_line_strength=config.min_line_strength) _plot_cross_section_panel(axes_flat[1], spectrum, xlim=xlim, show_components=show_components, component_label=component_label) if show_components: - _plot_secondary_component_panel( - axes_flat[2], - cia_dataset=cia_dataset, - spectrum=spectrum, - temperature_k=spectrum.temperature_k, - xlim=xlim, - ) + _plot_secondary_component_panel(axes_flat[2], cia_dataset=cia_dataset, spectrum=spectrum, temperature_k=spectrum.temperature_k, xlim=xlim) _plot_attenuation_panel(axes_flat[3], spectrum, xlim=xlim, show_components=True, component_label=component_label) _plot_transmittance_panel(axes_flat[4], transmittance, xlim=xlim, show_components=True, component_label=component_label) last_data_axis_index = 4 @@ -546,18 +415,15 @@ def _render_overview_page(fig, axes_flat, *, band, config, spectrum, transmittan ax.set_xlabel("Wavenumber [cm$^{-1}$]", fontsize=10) -def _compute_overview_products(args: argparse.Namespace): +def compute_overview_products(args: argparse.Namespace): if args.path_length_km <= 0.0: raise ValueError("path-length-km must be positive") band, config = _build_band_and_config(args) line_db = download_hitran_lines(config, band) line_list = load_hitran_line_list(config, band, line_db=line_db) - cia_dataset = _load_requested_cia_dataset(args, config) - _, _, spectrum, line_provider = _compute_requested_absorption_spectrum(args, line_db=line_db, cia_dataset=cia_dataset) - transmittance = compute_transmittance_spectrum( - spectrum=spectrum, - path_length_m=float(args.path_length_km) * 1000.0, - ) + cia_dataset = load_requested_cia_dataset(args, config) + _, _, spectrum, line_provider = compute_requested_absorption_spectrum(args, line_db=line_db, cia_dataset=cia_dataset) + transmittance = compute_transmittance_spectrum(spectrum=spectrum, path_length_m=float(args.path_length_km) * 1000.0) return band, config, line_list, spectrum, transmittance, cia_dataset, line_provider @@ -581,16 +447,14 @@ def __call__(self, parser, namespace, values, option_string=None) -> None: def build_molecule_overview_batch_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser( - description="Generate a combined multi-page PDF of molecule overview plots over one or more wavenumber ranges." - ) - parser.add_argument("--hitran-dir", type=Path, default=Path("hitran")) + parser = argparse.ArgumentParser(description="Generate a combined multi-page PDF of molecule overview plots over one or more wavenumber ranges.") + parser.add_argument("--hitran-dir", type=Path, default=default_hitran_dir()) parser.add_argument("--species", nargs="+", action=_SplitSpeciesAction, default=["H2", "CO2", "H2O", "CH4", "N2"]) parser.add_argument("--temperature-k", type=float, default=300.0) parser.add_argument("--pressure-bar", type=float, default=1.0) parser.add_argument("--resolution", type=float, default=1.0) parser.add_argument("--path-length-km", type=float, default=1.0) - parser.add_argument("--wn-range", dest="wn_ranges", action="append", type=_parse_wn_range, required=True) + parser.add_argument("--wn-range", dest="wn_ranges", action="append", type=parse_wn_range, required=True) parser.add_argument("--refresh-hitran", action="store_true") parser.add_argument("--broadening-composition", default=None) parser.add_argument("--refresh-cia", action="store_true") @@ -610,7 +474,7 @@ def run_overview_batch(args: argparse.Namespace) -> None: figure_path.parent.mkdir(parents=True, exist_ok=True) page_tasks: list[argparse.Namespace] = [] page_metadata: list[tuple[float, float, str, float, float]] = [] - for temperature_k, pressure_bar in _state_pairs(args): + for temperature_k, pressure_bar in state_pairs(args): for species in args.species: for wn_min, wn_max in args.wn_ranges: page_args = argparse.Namespace( @@ -639,31 +503,18 @@ def run_overview_batch(args: argparse.Namespace) -> None: ): band, config, line_list, spectrum, transmittance, cia_dataset, _ = result fig, axes = plt.subplots(nrows=5, ncols=1, figsize=(8.5, 11.0), squeeze=False) - _render_overview_page( - fig, - axes[:, 0], - band=band, - config=config, - spectrum=spectrum, - transmittance=transmittance, - line_list=line_list, - cia_dataset=cia_dataset, - ) + _render_overview_page(fig, axes[:, 0], band=band, config=config, spectrum=spectrum, transmittance=transmittance, line_list=line_list, cia_dataset=cia_dataset) pdf.savefig(fig) plt.close(fig) page_count += 1 component_label = _resolve_overview_component_label(cia_dataset=cia_dataset, spectrum=spectrum) if np.any(np.asarray(spectrum.sigma_cia_cm2_molecule, dtype=np.float64) > 0.0) else "none" - print( - f"Added page {page_count}: T={temperature_k:.1f} K | p={pressure_bar:.3f} bar | {species} | {wn_min:.3f}-{wn_max:.3f} cm^-1 | secondary={component_label}" - ) + print(f"Added page {page_count}: T={temperature_k:.1f} K | p={pressure_bar:.3f} bar | {species} | {wn_min:.3f}-{wn_max:.3f} cm^-1 | secondary={component_label}") print(f"Combined overview figure: {figure_path}") print(f"Pages: {page_count}") def build_overview_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser( - description="Generate a US-letter PDF overview of molecular lines, cross section, optional CIA, attenuation, and transmittance." - ) + parser = argparse.ArgumentParser(description="Generate a US-letter PDF overview of molecular lines, cross section, optional CIA, attenuation, and transmittance.") _add_common_arguments(parser) _add_external_cia_arguments(parser, required=False) parser.add_argument("--path-length-km", type=float, default=1.0) @@ -680,7 +531,7 @@ def run_overview(args: argparse.Namespace) -> None: figure_path = args.figure figure_path.parent.mkdir(parents=True, exist_ok=True) page_tasks: list[argparse.Namespace] = [] - for temperature_k, pressure_bar in _state_pairs(args): + for temperature_k, pressure_bar in state_pairs(args): page_args = argparse.Namespace(**vars(args)) page_args.temperature_k = float(temperature_k) page_args.pressure_bar = float(pressure_bar) @@ -689,23 +540,14 @@ def run_overview(args: argparse.Namespace) -> None: final_summary: tuple[object, object, object, object, str] | None = None with PdfPages(figure_path) as pdf: for temperature_k, pressure_bar, result in zip( - _selected_temperatures(args), - _selected_pressure_bars(args), + selected_temperatures(args), + selected_pressure_bars(args), _parallel_overview_page_products(page_tasks), strict=True, ): band, config, line_list, spectrum, transmittance, cia_dataset, broadening_summary = result fig, axes = plt.subplots(nrows=5, ncols=1, figsize=(8.5, 11.0), squeeze=False) - _render_overview_page( - fig, - axes[:, 0], - band=band, - config=config, - spectrum=spectrum, - transmittance=transmittance, - line_list=line_list, - cia_dataset=cia_dataset, - ) + _render_overview_page(fig, axes[:, 0], band=band, config=config, spectrum=spectrum, transmittance=transmittance, line_list=line_list, cia_dataset=cia_dataset) pdf.savefig(fig) plt.close(fig) print(f"Added page: T={temperature_k:.1f} K | p={pressure_bar:.3f} bar") @@ -714,17 +556,11 @@ def run_overview(args: argparse.Namespace) -> None: if final_summary is None: raise ValueError("at least one overview page is required") spectrum, transmittance, cia_dataset, _, broadening_summary = final_summary - _print_overview_summary( - figure_path=figure_path, - spectrum=spectrum, - transmittance=transmittance, - cia_dataset=cia_dataset, - broadening_summary=broadening_summary, - ) + _print_overview_summary(figure_path=figure_path, spectrum=spectrum, transmittance=transmittance, cia_dataset=cia_dataset, broadening_summary=broadening_summary) def _compute_overview_page_task(args: argparse.Namespace): - band, config, line_list, spectrum, transmittance, cia_dataset, line_provider = _compute_overview_products(args) + band, config, line_list, spectrum, transmittance, cia_dataset, line_provider = compute_overview_products(args) return band, config, line_list, spectrum, transmittance, cia_dataset, line_provider.broadening_summary() @@ -735,27 +571,44 @@ def _parallel_overview_page_products(tasks: list[argparse.Namespace]) -> Iterato return max_workers = min(len(tasks), os.cpu_count() or 1) ctx = process_pool_context() - with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as executor: - yield from executor.map(_compute_overview_page_task, tasks) + try: + with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as executor: + yield from executor.map(_compute_overview_page_task, tasks) + except PermissionError: + yield from _parallel_overview_page_products_via_subprocess(tasks, max_workers=max_workers) -def _selected_temperatures(args: argparse.Namespace) -> list[float]: - values = getattr(args, "temperature_k", [300.0]) - if isinstance(values, (list, tuple)): - return [float(value) for value in values] - return [float(values)] +def _subprocess_worker_entry(task_pickle_path: str, result_pickle_path: str) -> None: + with Path(task_pickle_path).open("rb") as handle: + task = pickle.load(handle) + result = _compute_overview_page_task(task) + with Path(result_pickle_path).open("wb") as handle: + pickle.dump(result, handle, protocol=pickle.HIGHEST_PROTOCOL) -def _selected_pressure_bars(args: argparse.Namespace) -> list[float]: - values = getattr(args, "pressure_bar", [1.0]) - if isinstance(values, (list, tuple)): - return [float(value) for value in values] - return [float(values)] +def _run_task_via_subprocess(task: argparse.Namespace): + code = ( + "from pyharp.spectra.hitran_molecule_plot import _subprocess_worker_entry; " + "import sys; " + "_subprocess_worker_entry(sys.argv[1], sys.argv[2])" + ) + with tempfile.TemporaryDirectory(prefix="pyharp_molecule_plot_task_") as tmpdir: + task_path = Path(tmpdir) / "task.pkl" + result_path = Path(tmpdir) / "result.pkl" + with task_path.open("wb") as handle: + pickle.dump(task, handle, protocol=pickle.HIGHEST_PROTOCOL) + subprocess.run( + [sys.executable, "-c", code, str(task_path), str(result_path)], + check=True, + ) + with result_path.open("rb") as handle: + return pickle.load(handle) -def _state_pairs(args: argparse.Namespace) -> list[tuple[float, float]]: - temperatures = _selected_temperatures(args) - pressures = _selected_pressure_bars(args) - if len(temperatures) != len(pressures): - raise ValueError("temperature_k and pressure_bar must have the same number of values") - return list(zip(temperatures, pressures, strict=True)) +def _parallel_overview_page_products_via_subprocess( + tasks: list[argparse.Namespace], + *, + max_workers: int, +) -> Iterator[tuple[object, object, object, object, object, object, str]]: + with ThreadPoolExecutor(max_workers=max_workers) as executor: + yield from executor.map(_run_task_via_subprocess, tasks) diff --git a/python/spectra/hitran_lines.py b/python/spectra/hitran_molecule_utils.py similarity index 69% rename from python/spectra/hitran_lines.py rename to python/spectra/hitran_molecule_utils.py index e6d4b5b..8c098bb 100644 --- a/python/spectra/hitran_lines.py +++ b/python/spectra/hitran_molecule_utils.py @@ -1,7 +1,8 @@ -"""HITRAN/HAPI line-download and line-opacity helpers.""" +"""Utility helpers for HITRAN molecular line workflows.""" from __future__ import annotations +import argparse import importlib import os from dataclasses import dataclass @@ -13,7 +14,9 @@ os.environ.setdefault("MPLCONFIGDIR", str(Path(tempfile.gettempdir()) / "spectra_matplotlib")) -from .config import SpectralBandConfig, SpectroscopyConfig +from .config import SpectralBandConfig, SpectroscopyConfig, parse_broadening_composition, resolve_hitran_cia_pair +from .hitran_cia_utils import load_cia_dataset +from .utils import build_band_from_range def _import_hapi(): @@ -23,6 +26,24 @@ def _import_hapi(): raise RuntimeError("HAPI is required for HITRAN line downloads.") from exc +def _resolve_continuum_sources(**kwargs): + from .spectrum import _resolve_continuum_sources as resolve_continuum_sources + + return resolve_continuum_sources(**kwargs) + + +def compute_absorption_spectrum_from_sources(**kwargs): + from .spectrum import compute_absorption_spectrum_from_sources as compute_from_sources + + return compute_from_sources(**kwargs) + + +def compute_transmittance_spectrum(**kwargs): + from .transmittance import compute_transmittance_spectrum as compute_transmittance + + return compute_transmittance(**kwargs) + + @dataclass(frozen=True) class LineDatabase: """Description of the cached HITRAN line table.""" @@ -296,39 +317,109 @@ def build_line_provider( ) -def plot_hitran_line_positions( - line_list: HitranLineList, - figure_path: Path, +def _build_band_and_config(args: argparse.Namespace) -> tuple[SpectralBandConfig, SpectroscopyConfig]: + band = build_band_from_range(tuple(args.wn_range), float(args.resolution)) + config = SpectroscopyConfig( + output_path=Path("output") / "unused.nc", + hitran_cache_dir=args.hitran_dir, + species_name=args.species, + broadening_composition=parse_broadening_composition(getattr(args, "broadening_composition", None)), + refresh_hitran=args.refresh_hitran, + ) + return band, config + + +def _resolve_cia_pair(args: argparse.Namespace, config: SpectroscopyConfig) -> str: + return str(args.cia_pair or f"{config.hitran_species.name}-{config.hitran_species.name}") + + +def _resolve_self_component_filename(args: argparse.Namespace, config: SpectroscopyConfig) -> str | None: + explicit = getattr(args, "cia_filename", None) + if explicit: + return str(explicit) + explicit_pair = getattr(args, "cia_pair", None) + if explicit_pair: + return resolve_hitran_cia_pair(explicit_pair).filename + return config.hitran_species.cia_filename + + +def load_requested_cia_dataset(args: argparse.Namespace, config: SpectroscopyConfig): + cia_filename = _resolve_self_component_filename(args, config) + if cia_filename is None: + return None + return load_cia_dataset( + cache_dir=args.hitran_dir, + filename=cia_filename, + pair=_resolve_cia_pair(args, config), + refresh=bool(args.refresh_cia), + ) + + +_MISSING = object() + + +def compute_requested_absorption_spectrum( + args: argparse.Namespace, *, - wavenumber_min_cm1: float | None = None, - wavenumber_max_cm1: float | None = None, - min_line_strength: float = 1.0e-27, -) -> None: - """Plot discrete HITRAN line positions and line intensities on log-log axes.""" - import matplotlib - - matplotlib.use("Agg") - import matplotlib.pyplot as plt - - figure_path.parent.mkdir(parents=True, exist_ok=True) - threshold = float(min_line_strength) - positive = line_list.line_intensity[line_list.line_intensity >= threshold] - if positive.size == 0: - raise ValueError("Line list does not contain any intensities above the minimum threshold.") - ymin = threshold - - fig, ax = plt.subplots(figsize=(12, 6)) - ax.vlines(line_list.wavenumber_cm1, ymin, line_list.line_intensity, color="0.2", alpha=0.2, linewidth=0.35) - ax.scatter(line_list.wavenumber_cm1, line_list.line_intensity, s=4.0, color="tab:blue", alpha=0.7, linewidths=0.0) - ax.set_xscale("log") - ax.set_yscale("log") - ax.set_ylim(ymin, max(float(np.max(positive)) * 1.2, ymin * 10.0)) - if wavenumber_min_cm1 is not None and wavenumber_max_cm1 is not None: - ax.set_xlim(float(wavenumber_min_cm1), float(wavenumber_max_cm1)) - ax.set_xlabel("Wavenumber [cm$^{-1}$]") - ax.set_ylabel("Line intensity $S$ [cm$^{-1}$/ (molecule cm$^{-2}$)]") - ax.set_title(f"{line_list.species_name} HITRAN line positions and intensities") - ax.grid(True, which="both", alpha=0.25) - fig.tight_layout() - fig.savefig(figure_path, dpi=150) - plt.close(fig) + line_db: LineDatabase | None = None, + cia_dataset=_MISSING, +): + band, config = _build_band_and_config(args) + temperature_k = float(args.temperature_k) + pressure_pa = float(args.pressure_bar) * 1.0e5 + if cia_dataset is _MISSING: + cia_dataset = load_requested_cia_dataset(args, config) + line_db = line_db or download_hitran_lines(config, band) + line_provider = build_line_provider(config, line_db) + grid = band.grid() + cia_cross_section_cm2_molecule = None + if cia_dataset is None: + cia_dataset, cia_cross_section_cm2_molecule = _resolve_continuum_sources( + config=config, + wavenumber_grid_cm1=grid, + temperature_k=temperature_k, + pressure_pa=pressure_pa, + ) + spectrum = compute_absorption_spectrum_from_sources( + species_name=config.hitran_species.name, + wavenumber_grid_cm1=grid, + temperature_k=temperature_k, + pressure_pa=pressure_pa, + line_provider=line_provider, + cia_dataset=cia_dataset, + cia_cross_section_cm2_molecule=cia_cross_section_cm2_molecule, + ) + return band, config, spectrum, line_provider + + +def compute_requested_transmittance(args: argparse.Namespace): + if args.path_length_km <= 0.0: + raise ValueError("path-length-km must be positive") + _, _, spectrum, line_provider = compute_requested_absorption_spectrum(args) + transmittance = compute_transmittance_spectrum( + spectrum=spectrum, + path_length_m=float(args.path_length_km) * 1000.0, + ) + return spectrum, transmittance, line_provider + + +def selected_temperatures(args: argparse.Namespace) -> list[float]: + values = getattr(args, "temperature_k", [300.0]) + if isinstance(values, (list, tuple)): + return [float(value) for value in values] + return [float(values)] + + +def selected_pressure_bars(args: argparse.Namespace) -> list[float]: + values = getattr(args, "pressure_bar", [1.0]) + if isinstance(values, (list, tuple)): + return [float(value) for value in values] + return [float(values)] + + +def state_pairs(args: argparse.Namespace) -> list[tuple[float, float]]: + temperatures = selected_temperatures(args) + pressures = selected_pressure_bars(args) + if len(temperatures) != len(pressures): + raise ValueError("temperature_k and pressure_bar must have the same number of values") + return list(zip(temperatures, pressures, strict=True)) diff --git a/python/spectra/output_names.py b/python/spectra/output_names.py deleted file mode 100644 index 44648d8..0000000 --- a/python/spectra/output_names.py +++ /dev/null @@ -1,55 +0,0 @@ -"""Default output filename helpers for spectroscopy CLIs.""" - -from __future__ import annotations - -from pathlib import Path - - -def _format_value(value: float | int | str, unit: str = "") -> str: - if isinstance(value, str): - text = value - else: - text = f"{float(value):g}" - return f"{text.replace('-', 'm').replace('.', 'p')}{unit}" - - -def _clean_token(value: object) -> str: - text = str(value).strip().lower() - pieces: list[str] = [] - previous_was_separator = False - for char in text: - if char.isalnum(): - pieces.append(char) - previous_was_separator = False - elif char == ".": - pieces.append("p") - previous_was_separator = False - elif not previous_was_separator: - pieces.append("_") - previous_was_separator = True - return "".join(pieces).strip("_") or "output" - - -def default_output_path( - *, - target_name: object, - plot_type: str, - temperature_k: float | str, - pressure_bar: float | str, - wn_range: tuple[float, float], - suffix: str, - output_dir: Path = Path("output"), -) -> Path: - """Return a default output path using the spectroscopy CLI filename pattern.""" - wn_min, wn_max = wn_range - stem = "_".join( - [ - _clean_token(target_name), - _clean_token(plot_type), - _format_value(temperature_k, "K"), - _format_value(pressure_bar, "bar"), - _format_value(wn_min), - _format_value(wn_max, "cm1"), - ] - ) - return output_dir / f"{stem}{suffix}" diff --git a/python/spectra/plot_cli.py b/python/spectra/plot_cli.py index 0e874ba..09cb32c 100644 --- a/python/spectra/plot_cli.py +++ b/python/spectra/plot_cli.py @@ -3,14 +3,17 @@ from __future__ import annotations import argparse -from concurrent.futures import ProcessPoolExecutor +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor import os from pathlib import Path +import pickle +import subprocess +import sys +import tempfile from textwrap import dedent -from . import atm_overview_cli, cia_plot_cli, molecule_plot_cli -from .output_names import _format_value, default_output_path -from .shared_cli import process_pool_context +from . import atm_overview, hitran_cia_plot, hitran_molecule_plot +from .utils import _format_value, default_hitran_dir, default_output_path, parse_wn_range, process_pool_context class _SplitSpeciesAction(argparse.Action): @@ -51,7 +54,7 @@ def _add_state_arguments(parser: argparse.ArgumentParser) -> None: def _add_common_arguments(parser: argparse.ArgumentParser, *, allow_multiple_ranges: bool = False) -> None: - parser.add_argument("--hitran-dir", type=Path, default=Path("hitran"), metavar="DIR", help="Directory for downloaded HITRAN line and CIA data.") + parser.add_argument("--hitran-dir", type=Path, default=default_hitran_dir(), metavar="DIR", help="Directory for downloaded HITRAN line and CIA data.") parser.add_argument("--output-dir", type=Path, default=Path("output"), metavar="DIR", help="Directory for auto-generated figure output paths.") parser.add_argument("--resolution", type=float, default=1.0, metavar="CM^-1", help="Wavenumber grid spacing in cm^-1.") parser.add_argument( @@ -65,14 +68,14 @@ def _add_common_arguments(parser: argparse.ArgumentParser, *, allow_multiple_ran "--wn-range", dest="wn_ranges", action="append", - type=molecule_plot_cli._parse_wn_range, + type=parse_wn_range, metavar="MIN,MAX", help="Wavenumber range in cm^-1. Repeat for multi-page overview PDFs.", ) else: parser.add_argument( "--wn-range", - type=molecule_plot_cli._parse_wn_range, + type=parse_wn_range, default=None, metavar="MIN,MAX", help="Wavenumber range in cm^-1. CIA plots default to 20,10000; molecular and mixture plots default to 20,2500.", @@ -349,41 +352,85 @@ def _parallel_plot_results( return [worker(task) for task in tasks] max_workers = min(len(tasks), os.cpu_count() or 1) ctx = process_pool_context() - with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as executor: - return list(executor.map(worker, tasks)) + try: + with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as executor: + return list(executor.map(worker, tasks)) + except PermissionError: + return _parallel_plot_results_via_subprocess(tasks, worker=worker, max_workers=max_workers) + + +def _subprocess_worker_entry(task_pickle_path: str, result_pickle_path: str, worker_name: str) -> None: + with Path(task_pickle_path).open("rb") as handle: + task = pickle.load(handle) + worker = globals()[worker_name] + result = worker(task) + with Path(result_pickle_path).open("wb") as handle: + pickle.dump(result, handle, protocol=pickle.HIGHEST_PROTOCOL) + + +def _run_task_via_subprocess(task: tuple[str, argparse.Namespace], *, worker_name: str) -> None: + code = ( + "from pyharp.spectra.plot_cli import _subprocess_worker_entry; " + "import sys; " + "_subprocess_worker_entry(sys.argv[1], sys.argv[2], sys.argv[3])" + ) + with tempfile.TemporaryDirectory(prefix="pyharp_plot_task_") as tmpdir: + task_path = Path(tmpdir) / "task.pkl" + result_path = Path(tmpdir) / "result.pkl" + with task_path.open("wb") as handle: + pickle.dump(task, handle, protocol=pickle.HIGHEST_PROTOCOL) + subprocess.run( + [sys.executable, "-c", code, str(task_path), str(result_path), worker_name], + check=True, + ) + with result_path.open("rb") as handle: + return pickle.load(handle) + + +def _parallel_plot_results_via_subprocess( + tasks: list[tuple[str, argparse.Namespace]], + *, + worker, + max_workers: int, +) -> list[None]: + worker_name = getattr(worker, "__name__", None) + if worker_name is None or globals().get(worker_name) is not worker: + raise ValueError("worker must be a module-level plot_cli function for subprocess fallback") + with ThreadPoolExecutor(max_workers=max_workers) as executor: + return list(executor.map(lambda task: _run_task_via_subprocess(task, worker_name=worker_name), tasks)) def _run_plot_task(task: tuple[str, argparse.Namespace]) -> None: task_kind, task_args = task if task_kind == "molecule_xsection": - molecule_plot_cli.run_xsection(task_args) + hitran_molecule_plot.run_xsection(task_args) return if task_kind == "molecule_attenuation": - molecule_plot_cli.run_attenuation(task_args) + hitran_molecule_plot.run_attenuation(task_args) return if task_kind == "molecule_transmission": - molecule_plot_cli.run_transmission(task_args) + hitran_molecule_plot.run_transmission(task_args) return if task_kind == "molecule_overview": - molecule_plot_cli.run_overview(task_args) + hitran_molecule_plot.run_overview(task_args) return if task_kind == "molecule_overview_batch": - molecule_plot_cli.run_overview_batch(task_args) + hitran_molecule_plot.run_overview_batch(task_args) return if task_kind == "cia_attenuation": - cia_plot_cli.run_attenuation(task_args) + hitran_cia_plot.run_attenuation(task_args) return if task_kind == "cia_transmission": - cia_plot_cli.run_transmission(task_args) + hitran_cia_plot.run_transmission(task_args) return if task_kind == "atm_attenuation": - atm_overview_cli.run_atm_attenuation(task_args, wn_range=task_args.wn_range) + atm_overview.run_atm_attenuation(task_args, wn_range=task_args.wn_range) return if task_kind == "atm_transmission": - atm_overview_cli.run_atm_transmission(task_args, wn_range=task_args.wn_range) + atm_overview.run_atm_transmission(task_args, wn_range=task_args.wn_range) return if task_kind == "atm_overview": - atm_overview_cli.run_atm_overview(task_args) + atm_overview.run_atm_overview(task_args) return raise ValueError(f"unsupported plot task: {task_kind}") @@ -526,7 +573,7 @@ def main(argv: list[str] | None = None) -> None: args = parser.parse_args(argv) if args.command == "binary": - cia_plot_cli.run_binary(_as_cia_args(args, default_pair="H2-H2", plot_type="binary")) + hitran_cia_plot.run_binary(_as_cia_args(args, default_pair="H2-H2", plot_type="binary")) return _validate_state_grid(args, parser) @@ -649,7 +696,7 @@ def main(argv: list[str] | None = None) -> None: suffix=".pdf", ) args.figure = args.figure or default_path - atm_overview_cli.run_atm_overview(_as_atm_overview_args(args, wn_ranges=wn_ranges)) + atm_overview.run_atm_overview(_as_atm_overview_args(args, wn_ranges=wn_ranges)) return species = args.species or ["H2O"] @@ -665,9 +712,9 @@ def main(argv: list[str] | None = None) -> None: ) args.figure = args.figure or default_path if len(species) == 1 and len(wn_ranges) == 1: - molecule_plot_cli.run_overview(_as_molecule_overview_args(args, species=species[0], wn_range=wn_ranges[0])) + hitran_molecule_plot.run_overview(_as_molecule_overview_args(args, species=species[0], wn_range=wn_ranges[0])) else: - molecule_plot_cli.run_overview_batch(_as_molecule_overview_batch_args(args, species=species, wn_ranges=wn_ranges)) + hitran_molecule_plot.run_overview_batch(_as_molecule_overview_batch_args(args, species=species, wn_ranges=wn_ranges)) return parser.error(f"unsupported plot command: {args.command}") diff --git a/python/spectra/spectrum.py b/python/spectra/spectrum.py index 59f35f9..623dfaa 100644 --- a/python/spectra/spectrum.py +++ b/python/spectra/spectrum.py @@ -16,9 +16,9 @@ import xarray as xr from .config import SpectroscopyConfig, SpectralBandConfig -from .dataset_io import write_dataset_via_tmp -from .hitran_cia import CiaDataset, download_cia_file, parse_cia_file -from .hitran_lines import HapiLineProvider, LineDatabase, build_line_provider, download_hitran_lines +from .dataset_io import write_dataset +from .hitran_cia_utils import CiaDataset, download_cia_file, parse_cia_file +from .hitran_molecule_utils import HapiLineProvider, LineDatabase, build_line_provider, download_hitran_lines from .mt_ckd_h2o import compute_mt_ckd_h2o_continuum_cross_section K_BOLTZMANN = 1.380649e-23 @@ -187,7 +187,7 @@ def write_spectrum_dataset(spectrum: AbsorptionSpectrum, output_path: Path) -> N output_path.parent.mkdir(parents=True, exist_ok=True) dataset = spectrum_to_dataset(spectrum) try: - write_dataset_via_tmp(dataset, output_path) + write_dataset(dataset, output_path) finally: dataset.close() diff --git a/python/spectra/transmittance.py b/python/spectra/transmittance.py index 8940e64..6764b3c 100644 --- a/python/spectra/transmittance.py +++ b/python/spectra/transmittance.py @@ -11,7 +11,7 @@ from .blackbody import compute_normalized_blackbody_curve from .config import SpectroscopyConfig, SpectralBandConfig -from .dataset_io import DEFAULT_NETCDF_ENGINE, write_dataset_via_tmp +from .dataset_io import DEFAULT_NETCDF_ENGINE, write_dataset from .spectrum import AbsorptionSpectrum, compute_absorption_spectrum @@ -91,7 +91,7 @@ def write_transmittance_dataset(transmittance: TransmittanceSpectrum, output_pat output_path.parent.mkdir(parents=True, exist_ok=True) dataset = transmittance_to_dataset(transmittance) try: - write_dataset_via_tmp(dataset, output_path, engine=DEFAULT_NETCDF_ENGINE) + write_dataset(dataset, output_path, engine=DEFAULT_NETCDF_ENGINE) finally: dataset.close() diff --git a/python/spectra/shared_cli.py b/python/spectra/utils.py similarity index 63% rename from python/spectra/shared_cli.py rename to python/spectra/utils.py index 69bbe68..321c9fe 100644 --- a/python/spectra/shared_cli.py +++ b/python/spectra/utils.py @@ -7,7 +7,7 @@ from pathlib import Path import sys -from .output_names import default_output_path as default_named_output_path +from .config import SpectralBandConfig class HelpFormatter(argparse.RawDescriptionHelpFormatter): @@ -29,17 +29,56 @@ def project_root() -> Path: return Path(__file__).resolve().parents[2] -def default_output_path() -> Path: - """Return the default NetCDF output path inside the project root.""" - return default_named_output_path( - target_name="CO2", - plot_type="xsection", - temperature_k=300.0, - pressure_bar=1.0, - wn_range=(20.0, 2500.0), - suffix=".nc", - output_dir=project_root() / "output", +def _format_value(value: float | int | str, unit: str = "") -> str: + if isinstance(value, str): + text = value + else: + text = f"{float(value):g}" + return f"{text.replace('-', 'm').replace('.', 'p')}{unit}" + + +def _clean_token(value: object) -> str: + text = str(value).strip().lower() + pieces: list[str] = [] + previous_was_separator = False + for char in text: + if char.isalnum(): + pieces.append(char) + previous_was_separator = False + elif char == ".": + pieces.append("p") + previous_was_separator = False + elif not previous_was_separator: + pieces.append("_") + previous_was_separator = True + return "".join(pieces).strip("_") or "output" + + +def default_output_path( + *, + target_name: object = "CO2", + plot_type: str = "xsection", + temperature_k: float | str = 300.0, + pressure_bar: float | str = 1.0, + wn_range: tuple[float, float] = (20.0, 2500.0), + suffix: str = ".nc", + output_dir: Path | None = None, +) -> Path: + """Return a default spectroscopy CLI output path.""" + if output_dir is None: + output_dir = project_root() / "output" + wn_min, wn_max = wn_range + stem = "_".join( + [ + _clean_token(target_name), + _clean_token(plot_type), + _format_value(temperature_k, "K"), + _format_value(pressure_bar, "bar"), + _format_value(wn_min), + _format_value(wn_max, "cm1"), + ] ) + return Path(output_dir) / f"{stem}{suffix}" def default_hitran_dir() -> Path: @@ -80,17 +119,34 @@ def parse_wn_range(value: str) -> tuple[float, float]: return lower, upper +def build_band_from_range( + wn_range: tuple[float, float], + resolution_cm1: float, + *, + name: str = "single_state", +) -> SpectralBandConfig: + """Build a validated regular spectral band from explicit bounds.""" + wn_min, wn_max = wn_range + if resolution_cm1 <= 0.0: + raise ValueError("resolution must be positive") + if wn_max < wn_min: + raise ValueError("wn-range max must be >= min") + return SpectralBandConfig(name, float(wn_min), float(wn_max), float(resolution_cm1)) + + def build_band(args: argparse.Namespace): """Build the standard single-state spectral band config.""" - from .config import SpectralBandConfig + return build_band_from_range(tuple(args.wn_range), float(args.resolution)) + - wn_min, wn_max = args.wn_range - return SpectralBandConfig("single_state", wn_min, wn_max, args.resolution) +def build_grid(args: argparse.Namespace, *, name: str = "single_state"): + """Build a validated spectral grid from CLI-style args.""" + return build_band_from_range(tuple(args.wn_range), float(args.resolution), name=name).grid() def default_cli_output_path(args: argparse.Namespace, *, suffix: str) -> Path: """Return the default named output path for a CLI command.""" - return default_named_output_path( + return default_output_path( target_name=args.species, plot_type=args.command, temperature_k=args.temperature_k, diff --git a/src/compound.hpp b/src/compound.hpp index de70e4a..d0662a5 100644 --- a/src/compound.hpp +++ b/src/compound.hpp @@ -2,8 +2,12 @@ // C/C++ #include +#include #include +// torch +#include + namespace harp { //! Map from string names to doubles. Used for defining species mole/mass @@ -24,4 +28,40 @@ Composition get_composition(const std::string& formula); */ double get_compound_weight(const Composition& composition); +//! \brief Normalize a composition map so its values sum to unity +/*! + * \param composition Unnormalized composition map + * \return A normalized composition map whose values sum to 1 + */ +inline Composition normalize_composition(const Composition& composition) { + double total = 0.0; + for (auto const& [name, value] : composition) total += value; + + TORCH_CHECK(total > 0.0, "Composition sum must be positive"); + + Composition normalized; + for (auto const& [name, value] : composition) { + normalized[name] = value / total; + } + return normalized; +} + +//! \brief Format a composition map as a comma-separated string +/*! + * Example output: ``H2:0.9,He:0.1`` + * + * \param composition Composition map + * \return Comma-separated composition string + */ +inline std::string format_composition(const Composition& composition) { + std::ostringstream oss; + bool first = true; + for (auto const& [name, value] : composition) { + if (!first) oss << ","; + first = false; + oss << name << ":" << value; + } + return oss.str(); +} + } // namespace harp diff --git a/src/opacity/mean_molecular_weight.hpp b/src/opacity/mean_molecular_weight.hpp deleted file mode 100644 index d00974a..0000000 --- a/src/opacity/mean_molecular_weight.hpp +++ /dev/null @@ -1,20 +0,0 @@ -#include - -namespace harp { - -extern std::vector species_weights; - -inline torch::Tensor mean_molecular_weight(torch::Tensor conc) { - // check if conc has the same number of dimensions as species_weights - TORCH_CHECK(species_weights.size() == conc.size(-1), - "The last dimension of 'conc' must be the same as number of " - "species defined in yaml"); - torch::Tensor ww = torch::tensor( - species_weights, - torch::TensorOptions().dtype(conc.dtype()).device(conc.device())); - - // dimension of conc is (ncol, nlyr, nspecies) - return (conc * ww.unsqueeze(0).squeeze(0)).sum(-1) / conc.sum(-1); -} - -} // namespace harp diff --git a/src/opacity/molecule_cia.cpp b/src/opacity/molecule_cia.cpp new file mode 100644 index 0000000..30c29e7 --- /dev/null +++ b/src/opacity/molecule_cia.cpp @@ -0,0 +1,155 @@ +// harp +#include "molecule_cia.hpp" + +#include +#include + +namespace harp { + +extern std::vector species_names; + +MoleculeCIAImpl::MoleculeCIAImpl(OpacityOptions const& options_) + : options(options_) { + TORCH_CHECK(options->opacity_files().size() == 1, + "Only one opacity file is allowed"); + TORCH_CHECK( + options->species_ids().size() == 1 || options->species_ids().size() == 2, + "MoleculeCIA expects one species for self-CIA or two species for binary " + "CIA"); + TORCH_CHECK(options->type().empty() || options->type() == "molecule-cia", + "Mismatch opacity type: ", options->type(), + " expecting 'molecule-cia'"); + reset(); +} + +void MoleculeCIAImpl::reset() { +#ifdef NETCDFOUTPUT + int fileid = open_file(options->opacity_files()[0]); + + int wavenumber_id = -1; + check_nc(nc_inq_varid(fileid, "wavenumber", &wavenumber_id), + "Missing required variable wavenumber"); + wavenumber = convert_wavenumber_to_cm1(read_1d_variable(fileid, "wavenumber"), + read_var_units(fileid, wavenumber_id), + "wavenumber"); + + int pressure_id = -1; + check_nc(nc_inq_varid(fileid, "pressure", &pressure_id), + "Missing required variable pressure"); + ln_pressure = + convert_pressure_to_pa(read_1d_variable(fileid, "pressure"), + read_var_units(fileid, pressure_id), "pressure") + .log(); + + int del_temp_id = -1; + check_nc(nc_inq_varid(fileid, "del_temperature", &del_temp_id), + "Missing required variable del_temperature"); + temperature_anomaly = convert_temperature_to_k( + read_1d_variable(fileid, "del_temperature"), + read_var_units(fileid, del_temp_id), "del_temperature"); + + int base_temp_id = -1; + check_nc(nc_inq_varid(fileid, "temperature", &base_temp_id), + "Missing required variable temperature"); + ln_temperature_base = convert_temperature_to_k( + read_1d_variable(fileid, "temperature"), + read_var_units(fileid, base_temp_id), "temperature") + .log() + .unsqueeze(-1); + + auto first = normalize_token(species_names.at(options->species_ids().at(0))); + auto second = first; + if (options->species_ids().size() == 2) { + second = normalize_token(species_names.at(options->species_ids().at(1))); + } + + std::vector candidates = {"binary_absorption_coefficient_" + + first + "_" + second}; + if (first != second) { + candidates.push_back("binary_absorption_coefficient_" + second + "_" + + first); + } + + int varid = -1; + std::string selected; + for (auto const& candidate : candidates) { + if (try_find_varid(fileid, candidate, &varid)) { + selected = candidate; + break; + } + } + TORCH_CHECK(!selected.empty(), "No CIA variable found for species pair ", + first, "-", second); + + auto sigma_binary = + convert_binary_cross_section_to_m5_per_mol2( + read_tensor_permuted(fileid, selected, + {"wavenumber", "pressure", "del_temperature"}), + read_var_units(fileid, varid), selected) + .unsqueeze(-1); + sigma_binary = apply_positive_fill(sigma_binary, selected); + ln_sigma_binary = sigma_binary.log(); + + check_nc(nc_close(fileid), "Failed to close NetCDF file"); +#else + TORCH_CHECK(false, "NetCDF support is not enabled"); +#endif + + register_buffer("wavenumber", wavenumber); + register_buffer("ln_pressure", ln_pressure); + register_buffer("temperature_anomaly", temperature_anomaly); + register_buffer("ln_sigma_binary", ln_sigma_binary); + register_buffer("ln_temperature_base", ln_temperature_base); +} + +torch::Tensor MoleculeCIAImpl::forward( + torch::Tensor conc, std::map const& kwargs) { + TORCH_CHECK(kwargs.count("pres") > 0, "pres is required in kwargs"); + TORCH_CHECK(kwargs.count("temp") > 0, "temp is required in kwargs"); + + auto const& pres = kwargs.at("pres"); + auto const& temp = kwargs.at("temp"); + + int const ncol = conc.size(0); + int const nlyr = conc.size(1); + TORCH_CHECK(pres.size(0) == ncol && pres.size(1) == nlyr, + "Invalid pres shape: ", pres.sizes()); + TORCH_CHECK(temp.size(0) == ncol && temp.size(1) == nlyr, + "Invalid temp shape: ", temp.sizes()); + + torch::Tensor wave_query; + if (kwargs.count("wavenumber") > 0) { + wave_query = kwargs.at("wavenumber"); + } else if (kwargs.count("wavelength") > 0) { + wave_query = 1.0e4 / kwargs.at("wavelength"); + } else { + wave_query = wavenumber; + } + TORCH_CHECK(wave_query.dim() == 1, + "MoleculeCIA expects a 1D wavenumber or wavelength grid"); + + auto lnp = pres.log(); + auto temperature_base = + interpn({lnp}, {ln_pressure}, ln_temperature_base).squeeze(-1).exp(); + auto del_temp = temp - temperature_base; + + int const nwave = wave_query.size(0); + auto wave = + wave_query.unsqueeze(-1).unsqueeze(-1).expand({nwave, ncol, nlyr}); + auto lnp_grid = lnp.unsqueeze(0).expand({nwave, ncol, nlyr}); + auto temp_grid = del_temp.unsqueeze(0).expand({nwave, ncol, nlyr}); + + auto coeff = interpn({wave, lnp_grid, temp_grid}, + {wavenumber, ln_pressure, temperature_anomaly}, + ln_sigma_binary, true) + .exp(); + + auto species0 = conc.select(-1, options->species_ids().at(0)); + auto species1 = conc.select(-1, options->species_ids().size() == 2 + ? options->species_ids().at(1) + : options->species_ids().at(0)); + return coeff * species0.unsqueeze(0).unsqueeze(-1) * + species1.unsqueeze(0).unsqueeze(-1); +} + +} // namespace harp diff --git a/src/opacity/molecule_cia.hpp b/src/opacity/molecule_cia.hpp new file mode 100644 index 0000000..0f3db86 --- /dev/null +++ b/src/opacity/molecule_cia.hpp @@ -0,0 +1,31 @@ +#pragma once + +// torch +#include +#include +#include +#include +#include + +// harp +#include "opacity_options.hpp" + +namespace harp { + +class MoleculeCIAImpl : public torch::nn::Cloneable { + public: + torch::Tensor wavenumber, ln_pressure, temperature_anomaly; + torch::Tensor ln_sigma_binary, ln_temperature_base; + + OpacityOptions options; + + MoleculeCIAImpl() : options(OpacityOptionsImpl::create()) {} + explicit MoleculeCIAImpl(OpacityOptions const& options_); + void reset() override; + + torch::Tensor forward(torch::Tensor conc, + std::map const& kwargs); +}; +TORCH_MODULE(MoleculeCIA); + +} // namespace harp diff --git a/src/opacity/molecule_line.cpp b/src/opacity/molecule_line.cpp new file mode 100644 index 0000000..4bbabef --- /dev/null +++ b/src/opacity/molecule_line.cpp @@ -0,0 +1,167 @@ +// harp +#include "molecule_line.hpp" + +#include + +#include +#include + +namespace harp { + +extern std::vector species_names; + +MoleculeLineImpl::MoleculeLineImpl(OpacityOptions const& options_) + : options(options_) { + TORCH_CHECK(options->opacity_files().size() == 1, + "Only one opacity file is allowed"); + + TORCH_CHECK(options->species_ids().size() == 1, + "Only one species is allowed"); + + TORCH_CHECK(options->species_ids()[0] >= 0, + "Invalid species_id: ", options->species_ids()[0]); + + TORCH_CHECK(options->type().empty() || options->type() == "molecule-line", + "Mismatch opacity type: ", options->type(), + " expecting 'molecule-line'"); + + reset(); +} + +void MoleculeLineImpl::reset() { +#ifdef NETCDFOUTPUT + int fileid = open_file(options->opacity_files()[0]); + + int wavenumber_id = -1; + check_nc(nc_inq_varid(fileid, "wavenumber", &wavenumber_id), + "Missing required variable wavenumber"); + wavenumber = convert_wavenumber_to_cm1(read_1d_variable(fileid, "wavenumber"), + read_var_units(fileid, wavenumber_id), + "wavenumber"); + + int pressure_id = -1; + check_nc(nc_inq_varid(fileid, "pressure", &pressure_id), + "Missing required variable pressure"); + ln_pressure = + convert_pressure_to_pa(read_1d_variable(fileid, "pressure"), + read_var_units(fileid, pressure_id), "pressure") + .log(); + + int del_temp_id = -1; + check_nc(nc_inq_varid(fileid, "del_temperature", &del_temp_id), + "Missing required variable del_temperature"); + temperature_anomaly = convert_temperature_to_k( + read_1d_variable(fileid, "del_temperature"), + read_var_units(fileid, del_temp_id), "del_temperature"); + + int base_temp_id = -1; + check_nc(nc_inq_varid(fileid, "temperature", &base_temp_id), + "Missing required variable temperature"); + ln_temperature_base = convert_temperature_to_k( + read_1d_variable(fileid, "temperature"), + read_var_units(fileid, base_temp_id), "temperature") + .log() + .unsqueeze(-1); + + auto const species_token = + normalize_token(species_names.at(options->species_ids().at(0))); + auto const line_name = "sigma_line_" + species_token; + + int line_varid = -1; + check_nc(nc_inq_varid(fileid, line_name.c_str(), &line_varid), + "Missing required variable " + line_name); + auto sigma_cross = + convert_line_cross_section_to_m2_per_mol( + read_tensor_permuted(fileid, line_name, + {"wavenumber", "pressure", "del_temperature"}), + read_var_units(fileid, line_varid), line_name) + .unsqueeze(-1); + + int nvars = 0; + check_nc(nc_inq_nvars(fileid, &nvars), "Failed to query variable count"); + auto const continuum_prefix = "sigma_continuum_" + species_token + "_"; + for (int i = 0; i < nvars; ++i) { + char name[NC_MAX_NAME + 1] = {}; + check_nc(nc_inq_varname(fileid, i, name), "Failed to query variable name"); + std::string varname(name); + if (varname.rfind(continuum_prefix, 0) != 0) continue; + + sigma_cross += + convert_line_cross_section_to_m2_per_mol( + read_tensor_permuted(fileid, varname, + {"wavenumber", "pressure", "del_temperature"}), + read_var_units(fileid, i), varname) + .unsqueeze(-1); + } + + sigma_cross = apply_positive_fill(sigma_cross, line_name); + ln_sigma_cross = sigma_cross.log(); + + check_nc(nc_close(fileid), "Failed to close NetCDF file"); +#else + TORCH_CHECK(false, "NetCDF support is not enabled"); +#endif + + // register all buffers + register_buffer("wavenumber", wavenumber); + register_buffer("ln_pressure", ln_pressure); + register_buffer("temperature_anomaly", temperature_anomaly); + register_buffer("ln_sigma_cross", ln_sigma_cross); + register_buffer("ln_temperature_base", ln_temperature_base); +} + +torch::Tensor MoleculeLineImpl::forward( + torch::Tensor conc, std::map const& kwargs) { + int ncol = conc.size(0); + int nlyr = conc.size(1); + + TORCH_CHECK(kwargs.count("pres") > 0, "pres is required in kwargs"); + TORCH_CHECK(kwargs.count("temp") > 0, "temp is required in kwargs"); + + auto const& pres = kwargs.at("pres"); + auto const& temp = kwargs.at("temp"); + + TORCH_CHECK(pres.size(0) == ncol && pres.size(1) == nlyr, + "Invalid pres shape: ", pres.sizes(), + "; needs to be (ncol, nlyr)"); + TORCH_CHECK(temp.size(0) == ncol && temp.size(1) == nlyr, + "Invalid temp shape: ", temp.sizes(), + "; needs to be (ncol, nlyr)"); + + torch::Tensor wave_query; + if (kwargs.count("wavenumber") > 0) { + wave_query = kwargs.at("wavenumber"); + } else if (kwargs.count("wavelength") > 0) { + wave_query = 1.0e4 / kwargs.at("wavelength"); + } else { + wave_query = wavenumber; + } + TORCH_CHECK(wave_query.dim() == 1, + "MoleculeLine expects a 1D wavenumber or wavelength grid"); + + auto lnp = pres.log(); + auto temperature_base = + interpn({lnp}, {ln_pressure}, ln_temperature_base).squeeze(-1).exp(); + auto tempa = temp - temperature_base; + + int const nwave = wave_query.size(0); + auto wave = + wave_query.unsqueeze(-1).unsqueeze(-1).expand({nwave, ncol, nlyr}); + lnp = lnp.unsqueeze(0).expand({nwave, ncol, nlyr}); + tempa = tempa.unsqueeze(0).expand({nwave, ncol, nlyr}); + + auto out = interpn({wave, lnp, tempa}, + {wavenumber, ln_pressure, temperature_anomaly}, + ln_sigma_cross, true) + .exp(); + + // Check species id in range + TORCH_CHECK(options->species_ids()[0] >= 0 && + options->species_ids()[0] < conc.size(2), + "Invalid species_id: ", options->species_ids()[0]); + + return out * + conc.select(-1, options->species_ids()[0]).unsqueeze(0).unsqueeze(-1); +} + +} // namespace harp diff --git a/src/opacity/rfm.hpp b/src/opacity/molecule_line.hpp similarity index 57% rename from src/opacity/rfm.hpp rename to src/opacity/molecule_line.hpp index 5bdde06..d08cf40 100644 --- a/src/opacity/rfm.hpp +++ b/src/opacity/molecule_line.hpp @@ -12,32 +12,32 @@ namespace harp { -class RFMImpl : public torch::nn::Cloneable { +class MoleculeLineImpl : public torch::nn::Cloneable { public: - //! data table coordinate axis - //! (nwave,) (npres,) (ntemp,) - torch::Tensor kwave, klnp, ktempa; + //! data table coordinate axes + //! (nwave,) (npres,) (ndeltemp,) + torch::Tensor wavenumber, ln_pressure, temperature_anomaly; - //! tabulated absorption x-section [ln(m^2/kmol)] - //! (nwave, npres, ntemp, 1) - torch::Tensor kdata; + //! tabulated ln(line + same-species continuum cross section) [ln(m^2/mol)] + //! (nwave, npres, ndeltemp, 1) + torch::Tensor ln_sigma_cross; - //! reference temperature profile + //! ln(base temperature profile) [ln(K)] //! (npres, 1) - torch::Tensor kreftem; + torch::Tensor ln_temperature_base; - //! options with which this `RFMImpl` was constructed + //! options with which this `MoleculeLineImpl` was constructed OpacityOptions options; //! Constructor to initialize the layer - RFMImpl() : options(OpacityOptionsImpl::create()) {} - explicit RFMImpl(OpacityOptions const& options_); + MoleculeLineImpl() : options(OpacityOptionsImpl::create()) {} + explicit MoleculeLineImpl(OpacityOptions const& options_); void reset() override; //! Get optical properties /*! * This function calculates the absorption x-section of the gas - * based on the tabulated data generated by the RFM. + * based on the tabulated data generated by the MoleculeLine. * * \param conc mole concentration [mol/m^3], (ncol, nlyr, nspecies) * @@ -50,6 +50,6 @@ class RFMImpl : public torch::nn::Cloneable { torch::Tensor forward(torch::Tensor conc, std::map const& kwargs); }; -TORCH_MODULE(RFM); +TORCH_MODULE(MoleculeLine); } // namespace harp diff --git a/src/opacity/nitrogen_cia.cpp_ b/src/opacity/nitrogen_cia.cpp_ deleted file mode 100644 index fe79535..0000000 --- a/src/opacity/nitrogen_cia.cpp_ +++ /dev/null @@ -1,131 +0,0 @@ -// C/C++ -#include -#include -#include -#include - -// canoe -#include -#include -#include - -// climath -#include - -// utils -#include - -// opacity -#include "nitrogen_cia.hpp" - -// netcdf -#ifdef NETCDFOUTPUT -extern "C" { -#include -} -#endif - -static int const n_rt_sets = 10; -static int const n_fd1_sets = 5; -static int const n_fd2_sets = 5; - -void N2N2CIA::LoadCoefficient(std::string fname, int) { - std::string line; - std::ifstream file(fname); - double value; - - // Roto-translational band - std::getline(file, line); - auto sline = Vectorize(line.c_str()); - int nwave = std::stoi(sline[3]), ntemp = n_rt_sets; - - rt_axis_.resize(ntemp + nwave); - rt_.resize(ntemp * nwave); - - for (int i = 0; i < ntemp; ++i) { - rt_axis_[i] = std::stod(sline[4]); - for (int j = 0; j < nwave; ++j) { - file >> value; - rt_axis_[ntemp + j] = value; - file >> value; - rt_[i * nwave + j] = std::max(value, 0.); - } - std::getline(file, line); - std::getline(file, line); - sline = Vectorize(line.c_str()); - } - rt_len_[0] = ntemp; - rt_len_[1] = nwave; - - // Fundamental band - nwave = std::stoi(sline[3]), ntemp = n_fd1_sets; - - fd1_axis_.resize(ntemp + nwave); - fd1_.resize(ntemp * nwave); - - for (int i = 0; i < ntemp; ++i) { - fd1_axis_[i] = std::stod(sline[4]); - for (int j = 0; j < nwave; ++j) { - file >> value; - fd1_axis_[ntemp + j] = value; - file >> value; - fd1_[i * nwave + j] = std::max(value, 0.); - } - std::getline(file, line); - std::getline(file, line); - sline = Vectorize(line.c_str()); - } - fd1_len_[0] = ntemp; - fd1_len_[1] = nwave; - - // Fundamental band - nwave = std::stoi(sline[3]), ntemp = n_fd2_sets; - - fd2_axis_.resize(ntemp + nwave); - fd2_.resize(ntemp * nwave); - - for (int i = 0; i < ntemp; ++i) { - fd2_axis_[i] = std::stod(sline[4]); - for (int j = 0; j < nwave; ++j) { - if (i == 2) { - file >> value; - fd2_axis_[ntemp + j] = value; - file >> value; - fd2_[i * nwave + j] = std::max(value, 0.); - } else { - file >> value; - fd2_axis_[ntemp + j] = value; - file >> value; - fd2_[i * nwave + j] = std::max(value, 0.); - file >> value; - } - } - std::getline(file, line); - std::getline(file, line); - sline = Vectorize(line.c_str()); - } - fd2_len_[0] = ntemp; - fd2_len_[1] = nwave; -} - -Real N2N2CIA::getAttenuation1(Real wave, AirParcel const& qfrac) const { - Real kk, temp = qfrac.w[IDN]; - - Real coor[2] = {temp, wave}; - - int iN2 = mySpeciesId(0); - - if (wave >= 0.02 && wave <= 554.) { - interpnf(&kk, coor, rt_.data(), rt_axis_.data(), rt_len_, 2); - } else if (wave >= 2000. && wave <= 2698.) { - interpnf(&kk, coor, fd2_.data(), fd2_axis_.data(), fd2_len_, 2); - } else if (wave >= 1850. && wave <= 3000.) - interpnf(&kk, coor, fd1_.data(), fd1_axis_.data(), fd1_len_, 2); - else - kk = 0.; - - Real num_n2 = qfrac.w[IPR] / (Constants::kBoltz * temp) * qfrac.w[iN2]; - - // cm^5 molecule^{-2} * (molecule m^{-3})^2 - return 1.E-10 * num_n2 * num_n2 * kk; -} diff --git a/src/opacity/nitrogen_cia.hpp b/src/opacity/nitrogen_cia.hpp deleted file mode 100644 index 55b495b..0000000 --- a/src/opacity/nitrogen_cia.hpp +++ /dev/null @@ -1,41 +0,0 @@ -#ifndef SRC_OPACITY_NITROGEN_CIA_HPP_ -#define SRC_OPACITY_NITROGEN_CIA_HPP_ - -// C/C++ -#include -#include - -// opacity -#include "absorber.hpp" - -class N2N2CIA : public Absorber { - public: - N2N2CIA() : Absorber("N2-N2") {} - - virtual ~N2N2CIA() {} - - void LoadCoefficient(std::string fname, int bid) override; - - Real GetAttenuation(Real wave1, Real wave2, - AirParcel const& var) const override { - Real k1 = getAttenuation1(wave1, var); - Real k2 = getAttenuation1(wave2, var); - return (k1 + k2) / 2.; - } - - protected: - Real getAttenuation1(Real wave, AirParcel const& var) const; - - size_t rt_len_[2]; - size_t fd1_len_[2]; - size_t fd2_len_[2]; - - std::vector rt_axis_; - std::vector fd1_axis_; - std::vector fd2_axis_; - std::vector rt_; - std::vector fd1_; - std::vector fd2_; -}; - -#endif // SRC_OPACITY_NITROGEN_CIA_HPP_ diff --git a/src/opacity/opacity_options.cpp b/src/opacity/opacity_options.cpp index cd6ce3e..18209c4 100644 --- a/src/opacity/opacity_options.cpp +++ b/src/opacity/opacity_options.cpp @@ -98,39 +98,4 @@ OpacityOptions OpacityOptionsImpl::from_yaml(std::string const& filename, return op; } -std::vector OpacityOptionsImpl::query_wavenumber() const { - if (type().compare(0, 3, "rfm") == 0) { - if (opacity_files().size() < 1) { - throw std::runtime_error("no opacity files specified for RFM opacity"); - } - return read_dimvar_netcdf(opacity_files()[0], "Wavenumber"); - } else if (type().compare(0, 9, "multiband") == 0) { - if (opacity_files().size() < 1) { - throw std::runtime_error( - "no opacity files specified for multiband opacity"); - } - return read_var_pt(opacity_files()[0], "wavenumber"); - } else { - throw std::runtime_error( - "unsupported opacity type for querying wavenumber"); - } -} - -std::vector OpacityOptionsImpl::query_weight() const { - if (type().compare(0, 3, "rfm") == 0) { - if (opacity_files().size() < 1) { - throw std::runtime_error("no opacity files specified for RFM opacity"); - } - return read_dimvar_netcdf(opacity_files()[0], "weights"); - } else if (type().compare(0, 9, "multiband") == 0) { - if (opacity_files().size() < 1) { - throw std::runtime_error( - "no opacity files specified for multiband opacity"); - } - return read_var_pt(opacity_files()[0], "weights"); - } else { - throw std::runtime_error("unsupported opacity type for querying weight"); - } -} - } // namespace harp diff --git a/src/opacity/opacity_options.hpp b/src/opacity/opacity_options.hpp index 881ca35..14538dc 100644 --- a/src/opacity/opacity_options.hpp +++ b/src/opacity/opacity_options.hpp @@ -35,9 +35,6 @@ struct OpacityOptionsImpl { os << "* verbose = " << (verbose() ? "true" : "false") << "\n"; } - std::vector query_wavenumber() const; - std::vector query_weight() const; - //! type of the opacity source ADD_ARG(std::string, type) = ""; diff --git a/src/opacity/oxygen_cia.cpp_ b/src/opacity/oxygen_cia.cpp_ deleted file mode 100644 index 0a9264f..0000000 --- a/src/opacity/oxygen_cia.cpp_ +++ /dev/null @@ -1,174 +0,0 @@ -// C/C++ -#include -#include -#include - -// canoe -#include -#include -#include - -// climath -#include - -// utils -#include - -// opacity -#include "oxygen_cia.hpp" - -// netcdf -#ifdef NETCDFOUTPUT -extern "C" { -#include -} -#endif - -static int const n_fd_sets = 15; -static int const n_a1dg_x3sg00_sets = 1; -static int const n_a1dg_x3sg10_sets = 1; -static int const n_ab_sets = 1; -static int const n_other_sets = 1; - -void O2O2CIA::LoadCoefficient(std::string fname, int) { - std::string line; - std::ifstream file(fname); - double value; - - // Fundamental band - std::getline(file, line); - auto sline = Vectorize(line.c_str()); - int nwave = std::stoi(sline[3]), ntemp = n_fd_sets; - - fd_axis_.resize(ntemp + nwave); - fd_.resize(ntemp * nwave); - - for (int i = 0; i < ntemp; ++i) { - fd_axis_[i] = std::stod(sline[4]); - for (int j = 0; j < nwave; ++j) { - file >> value; - fd_axis_[ntemp + j] = value; - file >> value; - fd_[i * nwave + j] = std::max(value, 0.); - } - std::getline(file, line); - std::getline(file, line); - sline = Vectorize(line.c_str()); - } - fd_len_[0] = ntemp; - fd_len_[1] = nwave; - - // a1dg_x3sg00 - nwave = std::stoi(sline[3]), ntemp = n_a1dg_x3sg00_sets; - - a1dg_x3sg00_axis_.resize(ntemp + nwave); - a1dg_x3sg00_.resize(ntemp * nwave); - - for (int i = 0; i < ntemp; ++i) { - a1dg_x3sg00_axis_[i] = std::stod(sline[4]); - for (int j = 0; j < nwave; ++j) { - file >> value; - a1dg_x3sg00_axis_[ntemp + j] = value; - file >> value; - a1dg_x3sg00_[i * nwave + j] = std::max(value, 0.); - } - std::getline(file, line); - std::getline(file, line); - sline = Vectorize(line.c_str()); - } - a1dg_x3sg00_len_[0] = ntemp; - a1dg_x3sg00_len_[1] = nwave; - - // a1dg_x3sg10 - nwave = std::stoi(sline[3]), ntemp = n_a1dg_x3sg10_sets; - - a1dg_x3sg10_axis_.resize(ntemp + nwave); - a1dg_x3sg10_.resize(ntemp * nwave); - - for (int i = 0; i < ntemp; ++i) { - a1dg_x3sg10_axis_[i] = std::stod(sline[4]); - for (int j = 0; j < nwave; ++j) { - file >> value; - a1dg_x3sg10_axis_[ntemp + j] = value; - file >> value; - a1dg_x3sg10_[i * nwave + j] = std::max(value, 0.); - } - std::getline(file, line); - std::getline(file, line); - sline = Vectorize(line.c_str()); - } - a1dg_x3sg10_len_[0] = ntemp; - a1dg_x3sg10_len_[1] = nwave; - - // A band - nwave = std::stoi(sline[3]), ntemp = n_ab_sets; - - ab_axis_.resize(ntemp + nwave); - ab_.resize(ntemp * nwave); - - for (int i = 0; i < ntemp; ++i) { - ab_axis_[i] = std::stod(sline[4]); - for (int j = 0; j < nwave; ++j) { - file >> value; - ab_axis_[ntemp + j] = value; - file >> value; - ab_[i * nwave + j] = std::max(value, 0.); - file >> value; - } - std::getline(file, line); - std::getline(file, line); - sline = Vectorize(line.c_str()); - } - ab_len_[0] = ntemp; - ab_len_[1] = nwave; - - // other bands - nwave = std::stoi(sline[3]), ntemp = n_other_sets; - - other_axis_.resize(ntemp + nwave); - other_.resize(ntemp * nwave); - - for (int i = 0; i < ntemp; ++i) { - other_axis_[i] = std::stod(sline[4]); - for (int j = 0; j < nwave; ++j) { - file >> value; - other_axis_[ntemp + j] = value; - file >> value; - other_[i * nwave + j] = std::max(value, 0.); - file >> value; - } - std::getline(file, line); - std::getline(file, line); - sline = Vectorize(line.c_str()); - } - other_len_[0] = ntemp; - other_len_[1] = nwave; -} - -Real O2O2CIA::getAttenuation1(Real wave, AirParcel const& qfrac) const { - Real kk, temp = qfrac.w[IDN]; - - Real coor[2] = {temp, wave}; - - int iO2 = mySpeciesId(0); - - if (wave >= 1150. && wave <= 1950.) { - interpnf(&kk, coor, fd_.data(), fd_axis_.data(), fd_len_, 2); - } else if (wave >= 7450. && wave <= 8487.) - interpnf(&kk, coor, a1dg_x3sg00_.data(), a1dg_x3sg00_axis_.data(), - a1dg_x3sg00_len_, 2); - else if (wave >= 9001. && wave <= 9997.) - interpnf(&kk, coor, a1dg_x3sg10_.data(), a1dg_x3sg10_axis_.data(), - a1dg_x3sg10_len_, 2); - else if (wave >= 12600. && wave <= 13839.) - interpnf(&kk, coor, ab_.data(), ab_axis_.data(), ab_len_, 2); - else if (wave >= 14996. && wave <= 29790.) - interpnf(&kk, coor, other_.data(), other_axis_.data(), other_len_, 2); - else - kk = 0.; - - Real num_o2 = qfrac.w[IPR] / (Constants::kBoltz * temp) * qfrac.w[iO2]; - - // cm^5 molecule^{-2} * (molecule m^{-3})^2 - return 1.E-10 * num_o2 * num_o2 * kk; -} diff --git a/src/opacity/oxygen_cia.hpp b/src/opacity/oxygen_cia.hpp deleted file mode 100644 index 0da3339..0000000 --- a/src/opacity/oxygen_cia.hpp +++ /dev/null @@ -1,46 +0,0 @@ -#ifndef SRC_OPACITY_OXYGEN_CIA_HPP_ -#define SRC_OPACITY_OXYGEN_CIA_HPP_ - -// C/C++ -#include - -// harp -#include "absorber.hpp" - -class O2O2CIA : public Absorber { - public: - O2O2CIA() : Absorber("O2-O2") {} - - virtual ~O2O2CIA() {} - - void LoadCoefficient(std::string fname, int bid) override; - - Real GetAttenuation(Real wave1, Real wave2, - AirParcel const& var) const override { - Real k1 = getAttenuation1(wave1, var); - Real k2 = getAttenuation1(wave2, var); - return (k1 + k2) / 2.; - } - - protected: - Real getAttenuation1(Real wave, AirParcel const& var) const; - - size_t fd_len_[2]; - size_t a1dg_x3sg00_len_[2]; - size_t a1dg_x3sg10_len_[2]; - size_t ab_len_[2]; - size_t other_len_[2]; - - std::vector fd_axis_; - std::vector a1dg_x3sg00_axis_; - std::vector a1dg_x3sg10_axis_; - std::vector ab_axis_; - std::vector other_axis_; - std::vector fd_; - std::vector a1dg_x3sg00_; - std::vector a1dg_x3sg10_; - std::vector ab_; - std::vector other_; -}; - -#endif // SRC_OPACITY_OXYGEN_CIA_HPP_ diff --git a/src/opacity/read_cia_ff.cpp b/src/opacity/read_cia_ff.cpp deleted file mode 100644 index 6be322f..0000000 --- a/src/opacity/read_cia_ff.cpp +++ /dev/null @@ -1,102 +0,0 @@ -// C/C++ -#include -#include -#include - -// opacity -#include "read_cia_ff.hpp" - -std::tuple read_cia_reform( - std::string filename) { - std::ifstream file{filename}; // open file - torch::Tensor data; // create storage array - std::vector temperature_axis; - std::vector spectral_axis; - - if (file.good()) { - int nx; // number of spectral points, horizontal - int ny; // number of temperature points, vertical - double temperature; // temprary storage for temperature data - double spectral; // temprary storage for spectral data - file >> ny >> nx; - for (int i = 0; i < ny; ++i) { - file >> temperature; - temperature_axis.push_back(temperature); - } - for (int i = 0; i < nx; ++i) { - file >> spectral; - spectral_axis.push_back(spectral); - } - - data = torch::empty({ny, nx}, torch::kFloat64); - for (int j = 0; j < ny; ++j) { - for (int i = 0; i < nx; ++i) { - double val; - file >> val; - data[j][i] = val; - } - } - } else { - throw std::runtime_error("Unable to open " + filename); - } - - return {data, torch::tensor(temperature_axis), - torch::tensor(spectral_axis)}; // return the data and axis as tensor -} - -// we are going to read the file twice, first to count # of rows and columns and -// second time -std::tuple read_freefree( - std::string filename) { - int num_of_row = 0; - int num_of_column = 1; - torch::Tensor data; // create storage array - std::vector temperature_axis; - std::vector spectral_axis; - double spectral; - double temperature; - std::ifstream cia_file{filename}; // open file - if (cia_file.good()) { - std::string line; - std::getline(cia_file, line); // get the first line - std::getline(cia_file, line); // get the second line - // calculate the # of columns base on the # of space - for (int j = 0; j < line.size(); ++j) { - if (line[j] == ' ') { - ++num_of_column; - } - } - while (std::getline(cia_file, line)) { - ++num_of_row; - } - } else { - throw std::runtime_error("Unable to open " + filename); - } - cia_file.close(); // close it - cia_file.open(filename); // open it again - if (cia_file.good()) { - std::string line1; // storage for the first line - std::getline(cia_file, line1); // Skip the first line - int ny = num_of_row; - int nx = num_of_column; - for (int j = 0; j < nx; ++j) { - cia_file >> temperature; - temperature_axis.push_back(temperature); // read off temperature axis - } - data = torch::empty({ny, nx}, torch::kFloat64); - for (int j = 0; j < ny; ++j) { - cia_file >> spectral; // skip first double and store it into spectal axis - spectral_axis.push_back(spectral); - for (int i = 0; i < nx; ++i) { - double val; - cia_file >> val; - data[j][i] = val; - } - } - } else { - throw std::runtime_error("Unable to open " + filename); - } - cia_file.close(); - return {data, torch::tensor(temperature_axis), - torch::tensor(spectral_axis)}; // return the data and axis as tensor -} diff --git a/src/opacity/read_cia_ff.hpp b/src/opacity/read_cia_ff.hpp deleted file mode 100644 index 152577c..0000000 --- a/src/opacity/read_cia_ff.hpp +++ /dev/null @@ -1,15 +0,0 @@ -// torch -#include - -//! The first array is the 2D data sheet of cross sections, the second vector is -//! the temperature axis, and the third is spectral axis - -//! read cia reform format file on temperature vs. spectral wavelength 2D data -//! set. -std::tuple read_cia_reform( - std::string filename); - -//! read free-free absorption format file on temperature vs. spectral wavelength -//! 2D data set. -std::tuple read_freefree( - std::string filename); diff --git a/src/opacity/read_rayleigh.cpp b/src/opacity/read_rayleigh.cpp deleted file mode 100644 index a52d2fd..0000000 --- a/src/opacity/read_rayleigh.cpp +++ /dev/null @@ -1,26 +0,0 @@ -// C/C++ -#include -#include -#include - -// opacity -#include "read_rayleigh.hpp" - -std::vector read_rayleigh(std::string file) { - std::ifstream data{file}; // open file - std::vector cross_section_output; // store the output - std::string line; // a temporary storage to get over the first line - // if stream OK = file readable - if (data.good()) { - double x; - std::getline(data, line); // Skip the first line - // as long as next value readable - while (data >> x) { - // put it into output - cross_section_output.push_back(x); - } - } else { - throw std::runtime_error("Unable to open " + file); - } - return cross_section_output; -} diff --git a/src/opacity/read_rayleigh.hpp b/src/opacity/read_rayleigh.hpp deleted file mode 100644 index 6069d8d..0000000 --- a/src/opacity/read_rayleigh.hpp +++ /dev/null @@ -1,6 +0,0 @@ -// C/C++ -#include -#include - -//! read input data for scattering cross-section and output as vectors -std::vector read_rayleigh(std::string file); diff --git a/src/opacity/rfm.cpp b/src/opacity/rfm.cpp deleted file mode 100644 index 46523c6..0000000 --- a/src/opacity/rfm.cpp +++ /dev/null @@ -1,171 +0,0 @@ -// base -#include - -// harp -#include -#include - -#include "rfm.hpp" - -// netcdf -#ifdef NETCDFOUTPUT -extern "C" { -#include -} -#endif - -namespace harp { - -extern std::vector species_names; - -RFMImpl::RFMImpl(OpacityOptions const& options_) : options(options_) { - TORCH_CHECK(options->opacity_files().size() == 1, - "Only one opacity file is allowed"); - - TORCH_CHECK(options->species_ids().size() == 1, - "Only one species is allowed"); - - TORCH_CHECK(options->species_ids()[0] >= 0, - "Invalid species_id: ", options->species_ids()[0]); - - TORCH_CHECK( - options->type().empty() || (options->type().compare(0, 3, "rfm") == 0), - "Mismatch opacity type: ", options->type()); - - reset(); -} - -void RFMImpl::reset() { - auto full_path = find_resource(options->opacity_files()[0]); - - // data table shape (nwave, npres, ntemp) - size_t kshape[3]; - -#ifdef NETCDFOUTPUT - int fileid, dimid, varid, err; - nc_open(full_path.c_str(), NC_NETCDF4, &fileid); - - err = nc_inq_dimid(fileid, "Wavenumber", &dimid); - TORCH_CHECK(err == NC_NOERR, nc_strerror(err)); - - err = nc_inq_dimlen(fileid, dimid, kshape); - TORCH_CHECK(err == NC_NOERR, nc_strerror(err)); - - err = nc_inq_dimid(fileid, "Pressure", &dimid); - TORCH_CHECK(err == NC_NOERR, nc_strerror(err)); - - err = nc_inq_dimlen(fileid, dimid, kshape + 1); - TORCH_CHECK(err == NC_NOERR, nc_strerror(err)); - - err = nc_inq_dimid(fileid, "TempGrid", &dimid); - TORCH_CHECK(err == NC_NOERR, nc_strerror(err)); - - err = nc_inq_dimlen(fileid, dimid, kshape + 2); - TORCH_CHECK(err == NC_NOERR, nc_strerror(err)); - - // wavenumber grid - kwave = torch::empty({(int)kshape[0]}, torch::kFloat64); - - // pressure grid - klnp = torch::empty({(int)kshape[1]}, torch::kFloat64); - - // temperatur grid - ktempa = torch::empty({(int)kshape[2]}, torch::kFloat64); - - // wave grid - err = nc_inq_varid(fileid, "Wavenumber", &varid); - TORCH_CHECK(err == NC_NOERR, nc_strerror(err)); - - err = nc_get_var_double(fileid, varid, kwave.data_ptr()); - TORCH_CHECK(err == NC_NOERR, nc_strerror(err)); - - // pressure grid - err = nc_inq_varid(fileid, "Pressure", &varid); - TORCH_CHECK(err == NC_NOERR, nc_strerror(err)); - - err = nc_get_var_double(fileid, varid, klnp.data_ptr()); - TORCH_CHECK(err == NC_NOERR, nc_strerror(err)); - // FIXME: dirty fix pressure unit - // klnp /= 100.; - - // change pressure to ln-pressure - klnp.log_(); - - // temperature grid - err = nc_inq_varid(fileid, "TempGrid", &varid); - TORCH_CHECK(err == NC_NOERR, nc_strerror(err)); - - err = nc_get_var_double(fileid, varid, ktempa.data_ptr()); - TORCH_CHECK(err == NC_NOERR, nc_strerror(err)); - - // reference temperature - kreftem = torch::empty({(int)kshape[1], 1}, torch::kFloat64); - err = nc_inq_varid(fileid, "Temperature", &varid); - TORCH_CHECK(err == NC_NOERR, nc_strerror(err)); - - err = nc_get_var_double(fileid, varid, kreftem.data_ptr()); - TORCH_CHECK(err == NC_NOERR, nc_strerror(err)); - - // data - kdata = torch::empty({(int)kshape[0], (int)kshape[1], (int)kshape[2], 1}, - torch::kFloat64); - auto name = species_names[options->species_ids()[0]]; - - err = nc_inq_varid(fileid, name.c_str(), &varid); - TORCH_CHECK(err == NC_NOERR, nc_strerror(err)); - - err = nc_get_var_double(fileid, varid, kdata.data_ptr()); - TORCH_CHECK(err == NC_NOERR, nc_strerror(err)); - - nc_close(fileid); -#endif - - // register all buffers - register_buffer("kwave", kwave); - register_buffer("klnp", klnp); - register_buffer("ktempa", ktempa); - register_buffer("kdata", kdata); - register_buffer("kreftem", kreftem); -} - -torch::Tensor RFMImpl::forward( - torch::Tensor conc, std::map const& kwargs) { - int nwave = kwave.size(0); - int ncol = conc.size(0); - int nlyr = conc.size(1); - - TORCH_CHECK(kwargs.count("pres") > 0, "pres is required in kwargs"); - TORCH_CHECK(kwargs.count("temp") > 0, "temp is required in kwargs"); - - auto const& pres = kwargs.at("pres"); - auto const& temp = kwargs.at("temp"); - - TORCH_CHECK(pres.size(0) == ncol && pres.size(1) == nlyr, - "Invalid pres shape: ", pres.sizes(), - "; needs to be (ncol, nlyr)"); - TORCH_CHECK(temp.size(0) == ncol && temp.size(1) == nlyr, - "Invalid temp shape: ", temp.sizes(), - "; needs to be (ncol, nlyr)"); - - // get temperature anomaly - auto lnp = pres.log(); - auto tempa = temp - interpn({lnp}, {klnp}, kreftem).squeeze(-1); - - // interpolate data - auto wave = kwave.unsqueeze(-1).unsqueeze(-1).expand({nwave, ncol, nlyr}); - lnp = lnp.unsqueeze(0).expand({nwave, ncol, nlyr}); - tempa = tempa.unsqueeze(0).expand({nwave, ncol, nlyr}); - - auto out = interpn({wave, lnp, tempa}, {kwave, klnp, ktempa}, kdata); - - // Check species id in range - TORCH_CHECK(options->species_ids()[0] >= 0 && - options->species_ids()[0] < conc.size(2), - "Invalid species_id: ", options->species_ids()[0]); - - // ln(m*2/kmol) -> 1/m - return 1.E-3 * out.exp() * - conc.select(-1, options->species_ids()[0]).unsqueeze(0).unsqueeze(-1); -} - -} // namespace harp diff --git a/src/radiation/radiation_band.cpp b/src/radiation/radiation_band.cpp index c9a01bf..3179535 100644 --- a/src/radiation/radiation_band.cpp +++ b/src/radiation/radiation_band.cpp @@ -8,9 +8,10 @@ #include #include #include +#include +#include #include #include -#include #include #include #include @@ -99,8 +100,8 @@ RadiationBandOptions RadiationBandOptionsImpl::from_yaml( op->disort() = create_disort_config(op->nwave(), op->ncol(), op->nlyr(), op->nstr()); op->disort()->header("running disort " + bd_name); - op->disort()->wave_lower(std::vector(op->nwave(), wmin)); - op->disort()->wave_upper(std::vector(op->nwave(), wmax)); + op->set_wave_lower(std::vector(op->nwave(), wmin)); + op->set_wave_upper(std::vector(op->nwave(), wmax)); if (band["flags"]) { op->disort()->flags(trim_copy(band["flags"].as())); } @@ -109,8 +110,8 @@ RadiationBandOptions RadiationBandOptionsImpl::from_yaml( } } else if (op->solver_name() == "toon") { op->toon() = ToonMcKay89OptionsImpl::create(); - op->toon()->wave_lower(std::vector(op->nwave(), wmin)); - op->toon()->wave_upper(std::vector(op->nwave(), wmax)); + op->set_wave_lower(std::vector(op->nwave(), wmin)); + op->set_wave_upper(std::vector(op->nwave(), wmax)); } else if (op->solver_name() == "twostr") { TORCH_CHECK(false, "twostr solver not implemented"); } else { @@ -139,12 +140,12 @@ void RadiationBandImpl::reset() { if (op->type() == "jit") { opacities[name] = torch::nn::AnyModule(JITOpacity(op)); nmax_prop = std::max(nmax_prop, 2 + op->nmom()); - } else if (op->type() == "rfm-lbl") { - auto a = RFM(op); + } else if (op->type() == "molecule-line") { + auto a = MoleculeLine(op); nmax_prop = std::max(nmax_prop, 1); opacities[name] = torch::nn::AnyModule(a); - } else if (op->type() == "rfm-ck") { - auto a = RFM(op); + } else if (op->type() == "molecule-cia") { + auto a = MoleculeCIA(op); nmax_prop = std::max(nmax_prop, 1); opacities[name] = torch::nn::AnyModule(a); } else if (op->type() == "multiband-ck") { diff --git a/src/radiation/radiation_band.hpp b/src/radiation/radiation_band.hpp index 142bde0..11403b6 100644 --- a/src/radiation/radiation_band.hpp +++ b/src/radiation/radiation_band.hpp @@ -31,8 +31,9 @@ using OpacityDict = std::map; * * The `RadiationBand` object recognizes the following opacity source types: * - "jit": user-defined opacity module - * - "rfm-lbl": line-by-line opacity defined on wavenumber grid - * - "rfm-ck": correlated-k opacity table computed from rfm line-by-line table + * - "molecule-line": line opacity from NetCDF cross-section dumps + * - "molecule-cia": CIA opacity from NetCDF binary absorption coefficient + * dumps * - "multiband-ck": multi-band correlated-k opacity table * - "wavetemp": opacity table defined on wavenumber and temperature grid (CIA) * - "fourcolumn": Four-column opacity table (aerosol) @@ -67,6 +68,36 @@ struct RadiationBandOptionsImpl { static std::shared_ptr from_yaml( std::string const& filename, std::string const& bd_name); + RadiationBandOptionsImpl& set_wave_lower(std::vector const& values) { + if (nwave() > 0) { + TORCH_CHECK(static_cast(values.size()) == nwave(), + "wave_lower size(", values.size(), + ") inconsistent with nwave(", nwave(), ")"); + } + + if (disort() != nullptr) disort()->wave_lower(values); + if (toon() != nullptr) toon()->wave_lower(values); + + TORCH_CHECK(disort() != nullptr || toon() != nullptr, + "Cannot set wave_lower before creating a solver"); + return *this; + } + + RadiationBandOptionsImpl& set_wave_upper(std::vector const& values) { + if (nwave() > 0) { + TORCH_CHECK(static_cast(values.size()) == nwave(), + "wave_upper size(", values.size(), + ") inconsistent with nwave(", nwave(), ")"); + } + + if (disort() != nullptr) disort()->wave_upper(values); + if (toon() != nullptr) toon()->wave_upper(values); + + TORCH_CHECK(disort() != nullptr || toon() != nullptr, + "Cannot set wave_upper before creating a solver"); + return *this; + } + std::shared_ptr clone() const { auto op = std::make_shared(*this); if (op->disort() != nullptr) { diff --git a/src/utils/mean_molecular_weight.hpp b/src/utils/mean_molecular_weight.hpp new file mode 100644 index 0000000..95a3f6e --- /dev/null +++ b/src/utils/mean_molecular_weight.hpp @@ -0,0 +1,28 @@ +#pragma once + +// C/C++ +#include + +// torch +#include + +namespace harp { + +extern std::vector species_weights; + +inline torch::Tensor mean_molecular_weight(torch::Tensor conc) { + TORCH_CHECK(species_weights.size() == static_cast(conc.size(-1)), + "The last dimension of 'conc' must match the number of species"); + + auto ww = torch::tensor( + species_weights, + torch::TensorOptions().dtype(conc.dtype()).device(conc.device())); + + int ndim = conc.dim(); + std::vector shape(ndim, 1); + shape.back() = static_cast(species_weights.size()); + + return (conc * ww.view(shape)).sum(-1) / conc.sum(-1); +} + +} // namespace harp diff --git a/src/utils/netcdf_opacity_utils.hpp b/src/utils/netcdf_opacity_utils.hpp new file mode 100644 index 0000000..a939152 --- /dev/null +++ b/src/utils/netcdf_opacity_utils.hpp @@ -0,0 +1,268 @@ +#pragma once + +// C/C++ +#include +#include +#include +#include +#include +#include +#include + +// base +#include + +// torch +#include + +// harp +#include + +#include + +// netcdf +#ifdef NETCDFOUTPUT +extern "C" { +#include +} +#endif + +namespace harp { + +inline std::string trim_copy(std::string value) { + auto const begin = value.find_first_not_of(" \t\r\n"); + if (begin == std::string::npos) return ""; + auto const end = value.find_last_not_of(" \t\r\n"); + return value.substr(begin, end - begin + 1); +} + +inline std::string lower_copy(std::string value) { + std::transform(value.begin(), value.end(), value.begin(), + [](unsigned char c) { return std::tolower(c); }); + return value; +} + +inline std::string normalize_token(std::string value) { + value = lower_copy(value); + + std::string out; + out.reserve(value.size()); + bool last_was_sep = false; + for (unsigned char c : value) { + if (std::isalnum(c)) { + out.push_back(static_cast(c)); + last_was_sep = false; + } else if (!last_was_sep) { + out.push_back('_'); + last_was_sep = true; + } + } + + while (!out.empty() && out.front() == '_') out.erase(out.begin()); + while (!out.empty() && out.back() == '_') out.pop_back(); + return out; +} + +#ifdef NETCDFOUTPUT +inline void check_nc(int status, std::string const& context) { + TORCH_CHECK(status == NC_NOERR, context, ": ", nc_strerror(status)); +} + +inline int open_file(std::string const& filename) { + int fileid = -1; + auto full_path = find_resource(filename); + check_nc(nc_open(full_path.c_str(), NC_NOWRITE, &fileid), + "Failed to open NetCDF file " + full_path); + return fileid; +} + +inline bool try_find_varid(int fileid, std::string const& varname, int* varid) { + auto const err = nc_inq_varid(fileid, varname.c_str(), varid); + if (err == NC_NOERR) return true; + TORCH_CHECK(err == NC_ENOTVAR, "Failed to query variable ", varname, ": ", + nc_strerror(err)); + return false; +} + +inline std::string read_var_units(int fileid, int varid) { + size_t len = 0; + auto const status = nc_inq_attlen(fileid, varid, "units", &len); + if (status == NC_ENOTATT) return ""; + check_nc(status, "Failed to read units attribute length"); + + std::string units(len, '\0'); + check_nc(nc_get_att_text(fileid, varid, "units", units.data()), + "Failed to read units attribute"); + return trim_copy(units); +} + +inline std::vector read_var_dim_names(int fileid, int varid) { + int ndims = 0; + check_nc(nc_inq_varndims(fileid, varid, &ndims), + "Failed to query variable rank"); + + std::vector dimids(ndims); + check_nc(nc_inq_vardimid(fileid, varid, dimids.data()), + "Failed to query variable dimensions"); + + std::vector names(ndims); + for (int i = 0; i < ndims; ++i) { + char dim_name[NC_MAX_NAME + 1] = {}; + check_nc(nc_inq_dimname(fileid, dimids[i], dim_name), + "Failed to query dimension name"); + names[i] = dim_name; + } + return names; +} + +inline torch::Tensor read_1d_variable(int fileid, std::string const& varname) { + int varid = -1; + check_nc(nc_inq_varid(fileid, varname.c_str(), &varid), + "Missing required variable " + varname); + + int ndims = 0; + check_nc(nc_inq_varndims(fileid, varid, &ndims), + "Failed to query variable rank"); + TORCH_CHECK(ndims == 1, "Variable ", varname, " must be 1D, got rank ", + ndims); + + size_t len = 0; + int dimid = -1; + check_nc(nc_inq_vardimid(fileid, varid, &dimid), + "Failed to query variable dimensions"); + check_nc(nc_inq_dimlen(fileid, dimid, &len), + "Failed to query variable length"); + + std::vector data(len); + check_nc(nc_get_var_double(fileid, varid, data.data()), + "Failed to read variable " + varname); + return torch::tensor(data, torch::kFloat64); +} + +inline torch::Tensor read_tensor_permuted( + int fileid, std::string const& varname, + std::vector const& target_dims) { + int varid = -1; + check_nc(nc_inq_varid(fileid, varname.c_str(), &varid), + "Missing required variable " + varname); + + auto const source_dims = read_var_dim_names(fileid, varid); + TORCH_CHECK(source_dims.size() == target_dims.size(), "Variable ", varname, + " has rank ", source_dims.size(), ", expected ", + target_dims.size()); + + std::vector dimids(source_dims.size()); + check_nc(nc_inq_vardimid(fileid, varid, dimids.data()), + "Failed to query variable dimensions"); + + std::vector shape(source_dims.size(), 0); + for (size_t i = 0; i < dimids.size(); ++i) { + size_t len = 0; + check_nc(nc_inq_dimlen(fileid, dimids[i], &len), + "Failed to query dimension length"); + shape[i] = static_cast(len); + } + + auto numel = std::accumulate(shape.begin(), shape.end(), int64_t{1}, + std::multiplies()); + std::vector data(numel); + check_nc(nc_get_var_double(fileid, varid, data.data()), + "Failed to read variable " + varname); + + auto tensor = torch::from_blob(data.data(), shape, torch::kFloat64).clone(); + + std::vector permute(target_dims.size(), -1); + for (size_t target = 0; target < target_dims.size(); ++target) { + auto it = + std::find(source_dims.begin(), source_dims.end(), target_dims[target]); + TORCH_CHECK(it != source_dims.end(), "Variable ", varname, + " is missing dimension ", target_dims[target]); + permute[target] = static_cast(it - source_dims.begin()); + } + + for (size_t i = 0; i < permute.size(); ++i) { + TORCH_CHECK(std::count(permute.begin(), permute.end(), permute[i]) == 1, + "Variable ", varname, + " dimensions cannot be mapped uniquely to target order"); + } + + return tensor.permute(permute).contiguous(); +} +#endif + +inline torch::Tensor convert_wavenumber_to_cm1(torch::Tensor values, + std::string const& units, + std::string const& varname) { + auto const u = lower_copy(trim_copy(units)); + if (u.empty() || u == "cm^-1" || u == "cm-1" || u == "1/cm") return values; + if (u == "m^-1" || u == "m-1" || u == "1/m") return values * 1.0e-2; + TORCH_CHECK(false, "Unsupported wavenumber units for ", varname, ": ", units); + return values; +} + +inline torch::Tensor convert_pressure_to_pa(torch::Tensor values, + std::string const& units, + std::string const& varname) { + auto const u = lower_copy(trim_copy(units)); + if (u.empty() || u == "pa") return values; + if (u == "bar") return values * 1.0e5; + if (u == "atm") return values * 101325.0; + TORCH_CHECK(false, "Unsupported pressure units for ", varname, ": ", units); + return values; +} + +inline torch::Tensor convert_temperature_to_k(torch::Tensor values, + std::string const& units, + std::string const& varname) { + auto const u = lower_copy(trim_copy(units)); + if (u.empty() || u == "k" || u == "kelvin") return values; + if (u == "c" || u == "degc" || u == "celsius") return values + 273.15; + TORCH_CHECK(false, "Unsupported temperature units for ", varname, ": ", + units); + return values; +} + +inline torch::Tensor apply_positive_fill(torch::Tensor values, + std::string const& quantity_name) { + auto const positive_mask = values > 0; + double fill_value = 1.0e-300; + if (positive_mask.any().item()) { + fill_value = values.masked_select(positive_mask).min().item(); + } + + TORCH_CHECK(std::isfinite(fill_value) && fill_value > 0.0, + "Invalid positive fill value for ", quantity_name, ": ", + fill_value); + + return torch::where(positive_mask, values, + torch::full_like(values, fill_value)); +} + +inline torch::Tensor convert_line_cross_section_to_m2_per_mol( + torch::Tensor values, std::string const& units, + std::string const& varname) { + auto const u = lower_copy(trim_copy(units)); + if (u == "cm^2 molecule^-1" || u == "cm^2/molecule") { + return values * (1.0e-4 * constants::Avogadro); + } + if (u == "m^2 mol^-1" || u == "m^2/mol") return values; + if (u == "m^2 kmol^-1" || u == "m^2/kmol") return values * 1.0e-3; + TORCH_CHECK(false, "Unsupported line/continuum cross-section units for ", + varname, ": ", units); + return values; +} + +inline torch::Tensor convert_binary_cross_section_to_m5_per_mol2( + torch::Tensor values, std::string const& units, + std::string const& varname) { + auto const u = lower_copy(trim_copy(units)); + if (u == "cm^5 molecule^-2" || u == "cm^5/molecule^2") { + return values * (1.0e-10 * constants::Avogadro * constants::Avogadro); + } + if (u == "m^5 mol^-2" || u == "m^5/mol^2") return values; + TORCH_CHECK(false, "Unsupported CIA binary coefficient units for ", varname, + ": ", units); + return values; +} + +} // namespace harp diff --git a/src/utils/write_column_profile.hpp b/src/utils/write_column_profile.hpp new file mode 100644 index 0000000..91b74ca --- /dev/null +++ b/src/utils/write_column_profile.hpp @@ -0,0 +1,142 @@ +#pragma once + +// C/C++ +#include +#include +#include +#include +#include +#include + +// torch +#include + +// harp +#include + +namespace harp { + +extern std::vector species_names; + +inline void write_column_profile(std::filesystem::path const& filename, + torch::Tensor const& dz, + torch::Tensor const& pres, + torch::Tensor const& temp, + torch::Tensor const& conc, + torch::Tensor const& band_flux) { + auto dz_1d = dz.to(torch::kCPU).to(torch::kFloat64).contiguous(); + auto pres_2d = pres.to(torch::kCPU).to(torch::kFloat64).contiguous(); + auto temp_2d = temp.to(torch::kCPU).to(torch::kFloat64).contiguous(); + auto conc_3d = conc.to(torch::kCPU).to(torch::kFloat64).contiguous(); + auto flux = band_flux.to(torch::kCPU).to(torch::kFloat64).contiguous(); + + TORCH_CHECK(dz_1d.dim() == 1, "dz must be 1D (nlyr)"); + TORCH_CHECK(pres_2d.dim() == 2, "pres must be 2D (ncol, nlyr)"); + TORCH_CHECK(temp_2d.dim() == 2, "temp must be 2D (ncol, nlyr)"); + TORCH_CHECK(conc_3d.dim() == 3, "conc must be 3D (ncol, nlyr, nspecies)"); + TORCH_CHECK(flux.dim() == 3, "band_flux must be 3D (ncol, nlyr+1, 2)"); + TORCH_CHECK(pres_2d.size(0) == 1 && temp_2d.size(0) == 1 && + conc_3d.size(0) == 1 && flux.size(0) == 1, + "write_column_profile currently expects ncol = 1"); + + int64_t const nlyr = dz_1d.size(0); + TORCH_CHECK(pres_2d.size(1) == nlyr, "pres size mismatch with dz"); + TORCH_CHECK(temp_2d.size(1) == nlyr, "temp size mismatch with dz"); + TORCH_CHECK(conc_3d.size(1) == nlyr, "conc size mismatch with dz"); + TORCH_CHECK(flux.size(1) == nlyr + 1, "band_flux must have nlyr+1 levels"); + TORCH_CHECK(flux.size(2) == 2, "band_flux last dimension must be 2"); + TORCH_CHECK(conc_3d.size(2) == static_cast(species_names.size()), + "conc species dimension must match species_names"); + + std::filesystem::create_directories(filename.parent_path()); + std::ofstream out(filename); + TORCH_CHECK(out, "Failed to open output file: ", filename.string()); + + auto mmw = mean_molecular_weight(conc_3d).squeeze(0).contiguous(); + auto total_conc = conc_3d.sum(-1).squeeze(0).contiguous(); + auto upward = flux.index({0, torch::indexing::Slice(), 0}).contiguous(); + auto downward = flux.index({0, torch::indexing::Slice(), 1}).contiguous(); + auto net = upward - downward; + int64_t const nspecies_out = conc_3d.size(2) - 1; + int const idx_width = 6; + int const value_width = 16; + + std::vector level_height_km(static_cast(nlyr), 0.0); + double cumulative_km = 0.0; + for (int64_t i = 0; i < nlyr; ++i) { + level_height_km[static_cast(i)] = cumulative_km; + cumulative_km += dz_1d[i].item() * 1.0e-3; + } + + out << "# IDX: layer index (1-based); final separated row is TOA boundary\n"; + out << "# HGT: bottom-of-layer level height [km]\n"; + out << "# PRE: layer pressure [bar]\n"; + out << "# TEM: layer temperature [K]\n"; + out << "# species fields: mole fraction [ppmv], first species omitted\n"; + out << "# MMW: mean molecular weight [g/mol]\n"; + out << "# FUP: upward flux [W/m^2]\n"; + out << "# FDN: downward flux [W/m^2]\n"; + out << "# FNT: net flux = FUP - FDN [W/m^2]\n"; + + out << std::setw(idx_width) << "IDX" << std::setw(value_width) << "HGT" + << std::setw(value_width) << "PRE" << std::setw(value_width) << "TEM"; + for (int64_t is = 1; is < conc_3d.size(2); ++is) { + auto const& name = species_names[static_cast(is)]; + out << std::setw( + std::max(value_width, static_cast(name.size()) + 2)) + << name; + } + out << std::setw(value_width) << "MMW" << std::setw(value_width) << "FUP" + << std::setw(value_width) << "FDN" << std::setw(value_width) << "FNT" + << "\n"; + + out << std::setprecision(8); + for (int64_t i = 0; i < nlyr; ++i) { + out << std::setw(idx_width) << (i + 1); + out << std::setw(value_width) << std::fixed + << level_height_km[static_cast(i)]; + out << std::setw(value_width) << std::defaultfloat + << pres_2d[0][i].item() * 1.0e-5; + out << std::setw(value_width) << temp_2d[0][i].item(); + + auto total = total_conc[i].item(); + for (int64_t is = 1; is < conc_3d.size(2); ++is) { + double ppmv = 0.0; + if (total > 0.0) { + ppmv = conc_3d[0][i][is].item() / total * 1.0e6; + } + auto const& name = species_names[static_cast(is)]; + out << std::setw( + std::max(value_width, static_cast(name.size()) + 2)) + << ppmv; + } + + out << std::setw(value_width) << mmw[i].item() * 1.0e3; + out << std::setw(value_width) << upward[i].item(); + out << std::setw(value_width) << downward[i].item(); + out << std::setw(value_width) << net[i].item() << "\n"; + } + + out << std::setw(idx_width) << (nlyr + 1); + out << std::setw(value_width) << std::fixed << cumulative_km; + out << std::setw(value_width) << std::defaultfloat + << (pres_2d[0][nlyr - 1].item() * 1.0e-5); + out << std::setw(value_width) << temp_2d[0][nlyr - 1].item(); + for (int64_t is = 1; is < conc_3d.size(2); ++is) { + double ppmv = 0.0; + auto total = total_conc[nlyr - 1].item(); + if (total > 0.0) { + ppmv = conc_3d[0][nlyr - 1][is].item() / total * 1.0e6; + } + auto const& name = species_names[static_cast(is)]; + out << std::setw( + std::max(value_width, static_cast(name.size()) + 2)) + << ppmv; + } + out << std::setw(value_width) << mmw[nlyr - 1].item() * 1.0e3; + out << std::setw(value_width) << upward[nlyr].item(); + out << std::setw(value_width) << downward[nlyr].item(); + out << std::setw(value_width) << net[nlyr].item() << "\n"; +} + +} // namespace harp diff --git a/src/utils/write_spectral_profile.hpp b/src/utils/write_spectral_profile.hpp new file mode 100644 index 0000000..75d0c95 --- /dev/null +++ b/src/utils/write_spectral_profile.hpp @@ -0,0 +1,73 @@ +#pragma once + +// C/C++ +#include +#include +#include +#include +#include + +// torch +#include + +namespace harp { + +inline void write_spectral_profile( + std::filesystem::path const& filename, + std::vector const& wavenumber, + std::vector> const& transmittance, + torch::Tensor const& flux_up_toa, torch::Tensor const& flux_down_boa) { + auto fup = flux_up_toa.to(torch::kCPU).to(torch::kFloat64).contiguous(); + auto fdn = flux_down_boa.to(torch::kCPU).to(torch::kFloat64).contiguous(); + + TORCH_CHECK(fup.dim() == 1, "flux_up_toa must be 1D"); + TORCH_CHECK(fdn.dim() == 1, "flux_down_boa must be 1D"); + TORCH_CHECK(static_cast(wavenumber.size()) == fup.size(0), + "wavenumber size mismatch with flux_up_toa"); + TORCH_CHECK(fup.size(0) == fdn.size(0), + "flux_up_toa and flux_down_boa size mismatch"); + + std::filesystem::create_directories(filename.parent_path()); + std::ofstream out(filename); + TORCH_CHECK(out, "Failed to open output file: ", filename.string()); + int const value_width = 16; + + out << "# WNO: wavenumber [cm^-1]\n"; + for (auto const& [name, values] : transmittance) { + out << "# " << name << ": total transmittance from " << name << " [-]\n"; + } + out << "# FUP: upward spectral flux at TOA [W/m^2/cm^-1]\n"; + out << "# FDN: downward spectral flux at BOA [W/m^2/cm^-1]\n"; + + out << std::setw(value_width) << "WNO"; + std::vector trans_cpu; + for (auto const& [name, values] : transmittance) { + auto tensor = values.to(torch::kCPU).to(torch::kFloat64).contiguous(); + TORCH_CHECK(tensor.dim() == 1, "transmittance tensor for ", name, + " must be 1D"); + TORCH_CHECK(tensor.size(0) == static_cast(wavenumber.size()), + "transmittance size mismatch for ", name); + trans_cpu.push_back(tensor); + out << std::setw( + std::max(value_width, static_cast(name.size()) + 2)) + << name; + } + out << std::setw(value_width) << "FUP" << std::setw(value_width) << "FDN" + << "\n"; + + out << std::setprecision(8); + for (size_t i = 0; i < wavenumber.size(); ++i) { + out << std::setw(value_width) << std::defaultfloat << wavenumber[i]; + for (size_t j = 0; j < trans_cpu.size(); ++j) { + auto const& name = transmittance[j].first; + out << std::setw( + std::max(value_width, static_cast(name.size()) + 2)) + << trans_cpu[j][static_cast(i)].item(); + } + out << std::setw(value_width) << fup[static_cast(i)].item(); + out << std::setw(value_width) << fdn[static_cast(i)].item() + << "\n"; + } +} + +} // namespace harp diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 319b0dc..fa38d87 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -4,7 +4,7 @@ enable_testing() -#setup_test(test_attenuator) +setup_test(test_attenuator) setup_test(test_interpn) setup_test(test_hypsometric) setup_test(test_trapez) diff --git a/tests/spectra/test_atm_overview_cli.py b/tests/spectra/test_atm_overview_cli.py index c7de40f..5a1cc45 100644 --- a/tests/spectra/test_atm_overview_cli.py +++ b/tests/spectra/test_atm_overview_cli.py @@ -2,24 +2,24 @@ import numpy as np -from pyharp.spectra.atm_overview_cli import ( +from pyharp.spectra.atm_overview import ( _find_binary_pairs, _parallel_mixture_overview_products, - _parse_composition, build_atm_overview_parser, compute_mixture_overview_products, + parse_composition, run_atm_overview, ) def test_parse_composition_normalizes_and_merges_duplicates() -> None: - composition = _parse_composition("H2O:1,H2:8,H2O:1") + composition = parse_composition("H2O:1,H2:8,H2O:1") assert composition == {"H2O": 0.2, "H2": 0.8} def test_parse_composition_accepts_cia_only_species() -> None: - composition = _parse_composition("H2:0.9,He:0.1") + composition = parse_composition("H2:0.9,He:0.1") assert composition == {"H2": 0.9, "He": 0.1} @@ -46,7 +46,7 @@ def test_atm_overview_parser_accepts_manifest_path_and_ranges(tmp_path) -> None: "--broadening-composition", "H2:0.85,He:0.15", "--wn-range=25,2500", - "--wn-range=2501,20000", + "--wn-range=2500,20000", "--manifest", str(tmp_path / "sources.json"), "--figure", @@ -56,7 +56,7 @@ def test_atm_overview_parser_accepts_manifest_path_and_ranges(tmp_path) -> None: assert args.composition == "H2O:0.1,H2:0.9" assert args.broadening_composition == "H2:0.85,He:0.15" - assert args.wn_ranges == [(25.0, 2500.0), (2501.0, 20000.0)] + assert args.wn_ranges == [(25.0, 2500.0), (2500.0, 20000.0)] assert args.manifest == tmp_path / "sources.json" assert args.figure == tmp_path / "overview.pdf" @@ -76,11 +76,11 @@ def test_compute_mixture_overview_reports_broadening_fallback(monkeypatch, tmp_p ) monkeypatch.setattr( - "pyharp.spectra.atm_overview_cli.download_hitran_lines", + "pyharp.spectra.atm_overview.download_hitran_lines", lambda config, band: type("LineDb", (), {"table_name": "co2_lines_20_22", "cache_dir": tmp_path})(), ) monkeypatch.setattr( - "pyharp.spectra.atm_overview_cli.load_hitran_line_list", + "pyharp.spectra.atm_overview.load_hitran_line_list", lambda config, band: type( "LineList", (), @@ -96,13 +96,13 @@ def cross_section_cm2_molecule(self, **kwargs): grid = np.asarray(kwargs["wavenumber_grid_cm1"], dtype=np.float64) return np.ones_like(grid) - monkeypatch.setattr("pyharp.spectra.atm_overview_cli.build_line_provider", lambda config, line_db: FakeLineProvider()) + monkeypatch.setattr("pyharp.spectra.atm_overview.build_line_provider", lambda config, line_db: FakeLineProvider()) monkeypatch.setattr( - "pyharp.spectra.atm_overview_cli.load_cia_dataset", + "pyharp.spectra.atm_overview.load_cia_dataset", lambda *args, **kwargs: type("Cia", (), {"source_path": tmp_path / "cia", "pair": "CO2-CO2", "interpolate_to_grid": lambda self, temperature_k, grid: np.zeros_like(grid)})(), ) monkeypatch.setattr( - "pyharp.spectra.atm_overview_cli.compute_mt_ckd_h2o_continuum_cross_section", + "pyharp.spectra.atm_overview.compute_mt_ckd_h2o_continuum_cross_section", lambda **kwargs: np.zeros_like(kwargs["wavenumber_grid_cm1"]), ) @@ -131,11 +131,11 @@ def test_compute_mixture_overview_weights_h2o_continuum_by_h2o_fraction(monkeypa ) monkeypatch.setattr( - "pyharp.spectra.atm_overview_cli.download_hitran_lines", + "pyharp.spectra.atm_overview.download_hitran_lines", lambda config, band: type("LineDb", (), {"table_name": f"{config.hitran_species.name.lower()}_lines", "cache_dir": tmp_path})(), ) monkeypatch.setattr( - "pyharp.spectra.atm_overview_cli.load_hitran_line_list", + "pyharp.spectra.atm_overview.load_hitran_line_list", lambda config, band: type( "LineList", (), @@ -151,13 +151,13 @@ def cross_section_cm2_molecule(self, **kwargs): grid = np.asarray(kwargs["wavenumber_grid_cm1"], dtype=np.float64) return np.zeros_like(grid) - monkeypatch.setattr("pyharp.spectra.atm_overview_cli.build_line_provider", lambda config, line_db: FakeLineProvider()) + monkeypatch.setattr("pyharp.spectra.atm_overview.build_line_provider", lambda config, line_db: FakeLineProvider()) monkeypatch.setattr( - "pyharp.spectra.atm_overview_cli.load_cia_dataset", + "pyharp.spectra.atm_overview.load_cia_dataset", lambda *args, **kwargs: type("Cia", (), {"source_path": tmp_path / "cia", "pair": "H2-H2", "interpolate_to_grid": lambda self, temperature_k, grid: np.zeros_like(grid)})(), ) monkeypatch.setattr( - "pyharp.spectra.atm_overview_cli.compute_mt_ckd_h2o_continuum_cross_section", + "pyharp.spectra.atm_overview.compute_mt_ckd_h2o_continuum_cross_section", lambda **kwargs: np.full_like(np.asarray(kwargs["wavenumber_grid_cm1"], dtype=np.float64), 10.0), ) @@ -224,16 +224,16 @@ def savefig(self, fig): )() monkeypatch.setattr( - "pyharp.spectra.atm_overview_cli._parallel_mixture_overview_products", + "pyharp.spectra.atm_overview._parallel_mixture_overview_products", lambda tasks: iter([products]), ) - monkeypatch.setattr("pyharp.spectra.atm_overview_cli._render_mixture_overview", lambda fig, axes, *, products: None) - monkeypatch.setattr("pyharp.spectra.atm_overview_cli.PdfPages", _DummyPdf) + monkeypatch.setattr("pyharp.spectra.atm_overview._render_mixture_overview", lambda fig, axes, *, products: None) + monkeypatch.setattr("pyharp.spectra.atm_overview.PdfPages", _DummyPdf) monkeypatch.setattr( - "pyharp.spectra.atm_overview_cli.plt.subplots", + "pyharp.spectra.atm_overview.plt.subplots", lambda **kwargs: (object(), np.array([[object()], [object()], [object()], [object()]])), ) - monkeypatch.setattr("pyharp.spectra.atm_overview_cli.plt.close", lambda fig: None) + monkeypatch.setattr("pyharp.spectra.atm_overview.plt.close", lambda fig: None) monkeypatch.setattr( "pathlib.Path.write_text", lambda self, text: written.setdefault(str(self), text), @@ -263,10 +263,10 @@ def __exit__(self, exc_type, exc, tb): def map(self, worker, tasks): return [worker(task) for task in tasks] - monkeypatch.setattr("pyharp.spectra.atm_overview_cli.process_pool_context", lambda: "ctx-token") - monkeypatch.setattr("pyharp.spectra.atm_overview_cli.ProcessPoolExecutor", DummyExecutor) + monkeypatch.setattr("pyharp.spectra.atm_overview.process_pool_context", lambda: "ctx-token") + monkeypatch.setattr("pyharp.spectra.atm_overview.ProcessPoolExecutor", DummyExecutor) monkeypatch.setattr( - "pyharp.spectra.atm_overview_cli._compute_mixture_overview_product_task", + "pyharp.spectra.atm_overview._compute_mixture_overview_product_task", lambda task: {"task": task}, ) @@ -283,4 +283,21 @@ def map(self, worker, tasks): {"task": ("args-1", (20.0, 25.0))}, {"task": ("args-2", (25.0, 30.0))}, ] + + +def test_parallel_mixture_overview_products_falls_back_to_subprocess_workers_when_process_pool_locks_are_unavailable(monkeypatch) -> None: + class FailingExecutor: + def __init__(self, *args, **kwargs): + raise PermissionError(13, "Permission denied") + + monkeypatch.setattr("pyharp.spectra.atm_overview.process_pool_context", lambda: "ctx-token") + monkeypatch.setattr("pyharp.spectra.atm_overview.ProcessPoolExecutor", FailingExecutor) + monkeypatch.setattr( + "pyharp.spectra.atm_overview._parallel_mixture_overview_products_via_subprocess", + lambda tasks, *, max_workers: iter([{"task": "subprocess"}]), + ) + + result = list(_parallel_mixture_overview_products([("args-1", (20.0, 25.0)), ("args-2", (25.0, 30.0))])) + + assert result == [{"task": "subprocess"}] assert created == {"max_workers": 2, "mp_context": "ctx-token"} diff --git a/tests/spectra/test_cia.py b/tests/spectra/test_cia.py index f2087e7..e07329a 100644 --- a/tests/spectra/test_cia.py +++ b/tests/spectra/test_cia.py @@ -2,7 +2,8 @@ import numpy as np -from pyharp.spectra.hitran_cia import compute_cia_attenuation_m1, compute_cia_transmission, download_cia_file, download_cia_file_by_name, find_cia_download_url, parse_cia_file, plot_cia_cross_section, plot_cia_attenuation_coefficient, plot_cia_transmission +from pyharp.spectra.hitran_cia_plot import plot_cia_attenuation_coefficient, plot_cia_cross_section, plot_cia_transmission +from pyharp.spectra.hitran_cia_utils import compute_cia_attenuation_m1, compute_cia_transmission, download_cia_file, download_cia_file_by_name, find_cia_download_url, parse_cia_file from pyharp.spectra.config import SpectroscopyConfig @@ -65,7 +66,7 @@ def test_interpolate_uses_only_temperatures_with_spectral_coverage(tmp_path: Pat def test_find_cia_download_url_prefers_matching_link(monkeypatch) -> None: html = 'file' - monkeypatch.setattr("pyharp.spectra.hitran_cia._download_text", lambda _: html) + monkeypatch.setattr("pyharp.spectra.hitran_cia_utils._download_text", lambda _: html) url = find_cia_download_url("https://hitran.org/cia/", "CO2-CO2_2024.cia") assert url == "https://hitran.org/files/CO2-CO2_2024.cia" @@ -82,7 +83,7 @@ def _fail(*args, **kwargs): called["value"] = True raise AssertionError("network should not be used when cache exists") - monkeypatch.setattr("pyharp.spectra.hitran_cia.urlopen", _fail) + monkeypatch.setattr("pyharp.spectra.hitran_cia_utils.urlopen", _fail) path = download_cia_file(config) assert path == target assert called["value"] is False @@ -98,7 +99,7 @@ def _fail(*args, **kwargs): called["value"] = True raise AssertionError("network should not be used when cache exists") - monkeypatch.setattr("pyharp.spectra.hitran_cia.urlopen", _fail) + monkeypatch.setattr("pyharp.spectra.hitran_cia_utils.urlopen", _fail) path = download_cia_file_by_name(cache_dir=tmp_path, filename="H2-H2_2011.cia") assert path == target assert called["value"] is False diff --git a/tests/spectra/test_cli.py b/tests/spectra/test_cli.py index 3e1caab..bf1be48 100644 --- a/tests/spectra/test_cli.py +++ b/tests/spectra/test_cli.py @@ -5,9 +5,9 @@ import xarray as xr -from pyharp.spectra.dataset_io import combine_band_datasets, write_dataset_via_tmp -from pyharp.spectra.dump_cli import _args_for_wn_range, _composition_transmission_dataset, _composition_xsection_dataset, _output_path_for_wn_range, _pair_xsection_dataset, _resolve_pair_filename, _species_transmission_dataset, _stack_state_grid_datasets, _xsection_dataset, build_parser, main -from pyharp.spectra.shared_cli import default_hitran_dir, default_orton_xiz_cia_dir, default_output_path, project_root +from pyharp.spectra.dataset_io import combine_band_datasets, write_dataset +from pyharp.spectra.dump_cli import _args_for_wn_range, _composition_transmission_dataset, _composition_xsection_dataset, _compute_xsection_band, _output_path_for_wn_range, _pair_xsection_dataset, _parallel_band_results, _resolve_pair_filename, _species_transmission_dataset, _stack_state_grid_datasets, _xsection_dataset, build_parser, main +from pyharp.spectra.utils import default_hitran_dir, default_orton_xiz_cia_dir, default_output_path, project_root from pyharp.spectra.spectrum import AbsorptionSpectrum @@ -189,7 +189,7 @@ def interpolate_to_grid(self, *, temperature_k, wavenumber_grid_cm1): dataset = _pair_xsection_dataset(_args_for_wn_range(args, args.wn_ranges[0])) try: assert dataset.attrs["pair_name"] == "H2-H2" - assert np.allclose(dataset["binary_absorption_coefficient"].values, np.array([4.0e-46, 4.0e-46])) + assert np.allclose(dataset["binary_absorption_coefficient"].values, np.array([4.0e-46, 4.0e-46, 4.0e-46])) finally: dataset.close() @@ -432,7 +432,7 @@ def test_cli_xsection_reports_broadening_summary(monkeypatch, tmp_path, capsys) "requested=h2:0.900,he:0.100 -> effective=air:1.000 (fallback: h2->air, he->air)", ), ) - monkeypatch.setattr("pyharp.spectra.dump_cli.write_dataset_via_tmp", lambda dataset, output_path, *, engine: None) + monkeypatch.setattr("pyharp.spectra.dump_cli.write_dataset", lambda dataset, output_path, *, engine: None) main() @@ -472,7 +472,7 @@ def test_cli_xsection_writes_one_file_per_range(monkeypatch, tmp_path, capsys) - lambda tasks, *, worker: [worker(task) for task in tasks], ) monkeypatch.setattr( - "pyharp.spectra.dump_cli.write_dataset_via_tmp", + "pyharp.spectra.dump_cli.write_dataset", lambda dataset, output_path, *, engine: written.append((tuple(dataset["wavenumber"].values), output_path, engine)), ) @@ -523,7 +523,7 @@ def test_cli_xsection_composition_writes_one_file_with_component_fields(monkeypa lambda tasks, *, worker: [worker(task) for task in tasks], ) monkeypatch.setattr( - "pyharp.spectra.dump_cli.write_dataset_via_tmp", + "pyharp.spectra.dump_cli.write_dataset", lambda dataset, output_path, *, engine: written.append((sorted(dataset.data_vars), dataset["sigma_total"].dims, tuple(dataset["del_temperature"].values), tuple(dataset["pressure"].values), tuple(dataset["temperature"].values), output_path, engine)), ) @@ -571,7 +571,7 @@ def test_cli_xsection_uses_output_dir_for_generated_path(monkeypatch, tmp_path, ), ) monkeypatch.setattr( - "pyharp.spectra.dump_cli.write_dataset_via_tmp", + "pyharp.spectra.dump_cli.write_dataset", lambda dataset, output_path, *, engine: written.append(output_path), ) @@ -665,7 +665,7 @@ def test_cli_xsection_temperature_list_writes_temperature_stacked_dataset(monkey ) written = [] monkeypatch.setattr( - "pyharp.spectra.dump_cli.write_dataset_via_tmp", + "pyharp.spectra.dump_cli.write_dataset", lambda dataset, output_path, *, engine: written.append((dataset["sigma_total"].dims, tuple(dataset["del_temperature"].values), tuple(dataset["pressure"].values), tuple(dataset["temperature"].values), output_path, engine)), ) @@ -717,7 +717,7 @@ def fake_parallel(tasks, *, worker): ] monkeypatch.setattr("pyharp.spectra.dump_cli._parallel_band_results", fake_parallel) - monkeypatch.setattr("pyharp.spectra.dump_cli.write_dataset_via_tmp", lambda dataset, output_path, *, engine: None) + monkeypatch.setattr("pyharp.spectra.dump_cli.write_dataset", lambda dataset, output_path, *, engine: None) main() @@ -767,17 +767,55 @@ def test_combine_band_datasets_preserves_band_metadata() -> None: combined.close() -def test_write_dataset_via_tmp_creates_missing_output_parent(tmp_path) -> None: +def test_write_dataset_creates_missing_output_parent(tmp_path) -> None: dataset = xr.Dataset(coords={"wavenumber": ("wavenumber", np.array([1.0, 2.0]))}, data_vars={"sigma_total": ("wavenumber", np.array([3.0, 4.0]))}) output_path = tmp_path / "nested" / "pair.nc" try: - write_dataset_via_tmp(dataset, output_path, engine="scipy") + write_dataset(dataset, output_path, engine="scipy") assert output_path.exists() finally: dataset.close() +def test_write_dataset_disables_hdf5_file_locking(monkeypatch, tmp_path) -> None: + dataset = xr.Dataset(coords={"wavenumber": ("wavenumber", np.array([1.0, 2.0]))}, data_vars={"sigma_total": ("wavenumber", np.array([3.0, 4.0]))}) + output_path = tmp_path / "nested" / "pair.nc" + captured = {} + + def fake_to_netcdf(self, path, *, engine): + captured["path"] = path + captured["engine"] = engine + captured["locking"] = __import__("os").environ.get("HDF5_USE_FILE_LOCKING") + + monkeypatch.setattr(xr.Dataset, "to_netcdf", fake_to_netcdf) + + try: + write_dataset(dataset, output_path, engine="scipy") + assert captured == {"path": output_path, "engine": "scipy", "locking": "FALSE"} + finally: + dataset.close() + + +def test_parallel_band_results_falls_back_to_subprocess_workers_when_process_pool_locks_are_unavailable(monkeypatch) -> None: + tasks = [("species", "first"), ("pair", "second")] + + class FailingExecutor: + def __init__(self, *args, **kwargs): + raise PermissionError(13, "Permission denied") + + monkeypatch.setattr("pyharp.spectra.dump_cli.process_pool_context", lambda: "ctx-token") + monkeypatch.setattr("pyharp.spectra.dump_cli.ProcessPoolExecutor", FailingExecutor) + monkeypatch.setattr( + "pyharp.spectra.dump_cli._parallel_band_results_via_subprocess", + lambda task_items, *, worker, max_workers: [("subprocess-result", None)], + ) + + results = _parallel_band_results(tasks, worker=_compute_xsection_band) + + assert results == [("subprocess-result", None)] + + def test_xsection_dataset_keeps_only_sigma_fields() -> None: spectrum = AbsorptionSpectrum( species_name="H2O", @@ -988,7 +1026,7 @@ def test_cli_xsection_composition_multi_range_uses_composition_worker(monkeypatc lambda tasks, *, worker: [worker(task) for task in tasks], ) monkeypatch.setattr( - "pyharp.spectra.dump_cli.write_dataset_via_tmp", + "pyharp.spectra.dump_cli.write_dataset", lambda dataset, output_path, *, engine: written.append((tuple(dataset["wavenumber"].values), output_path, engine)), ) diff --git a/tests/spectra/test_config.py b/tests/spectra/test_config.py index 8a5b503..c885f6d 100644 --- a/tests/spectra/test_config.py +++ b/tests/spectra/test_config.py @@ -6,6 +6,13 @@ resolve_hitran_cia_pair, resolve_hitran_species, ) +import numpy as np + + +def test_spectral_band_grid_includes_both_endpoints() -> None: + band = SpectralBandConfig("single_state", 20.0, 22.0, 1.0) + + assert np.allclose(band.grid(), np.array([20.0, 21.0, 22.0])) def test_resolve_hitran_species_is_case_insensitive() -> None: diff --git a/tests/spectra/test_hitran_lines.py b/tests/spectra/test_hitran_lines.py index 4756265..6c509f8 100644 --- a/tests/spectra/test_hitran_lines.py +++ b/tests/spectra/test_hitran_lines.py @@ -1,12 +1,12 @@ import numpy as np from pyharp.spectra.config import SpectralBandConfig, SpectroscopyConfig -from pyharp.spectra.hitran_lines import ( +from pyharp.spectra.hitran_molecule_plot import plot_hitran_line_positions +from pyharp.spectra.hitran_molecule_utils import ( HapiLineProvider, HitranLineList, build_line_provider, download_hitran_lines, load_hitran_line_list, - plot_hitran_line_positions, ) @@ -45,7 +45,7 @@ def absorptionCoefficient_Voigt(self, **kwargs): def test_download_hitran_lines_uses_band_bounds(monkeypatch, tmp_path) -> None: fake = FakeHapi() - monkeypatch.setattr("pyharp.spectra.hitran_lines._import_hapi", lambda: fake) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils._import_hapi", lambda: fake) config = SpectroscopyConfig(output_path=tmp_path / "out.nc", hitran_cache_dir=tmp_path / "cache") band = SpectralBandConfig("single_state", 25.0, 2500.0, 1.0) @@ -65,7 +65,7 @@ def test_download_hitran_lines_uses_band_bounds(monkeypatch, tmp_path) -> None: def test_download_hitran_lines_skips_fetch_when_cache_exists(monkeypatch, tmp_path) -> None: fake = FakeHapi() - monkeypatch.setattr("pyharp.spectra.hitran_lines._import_hapi", lambda: fake) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils._import_hapi", lambda: fake) config = SpectroscopyConfig(output_path=tmp_path / "out.nc", hitran_cache_dir=tmp_path / "cache") band = SpectralBandConfig("single_state", 25.0, 2500.0, 1.0) config.hitran_cache_dir.mkdir(parents=True, exist_ok=True) @@ -82,7 +82,7 @@ def test_download_hitran_lines_skips_fetch_when_cache_exists(monkeypatch, tmp_pa def test_download_hitran_lines_refetches_contaminated_cache(monkeypatch, tmp_path) -> None: fake = FakeHapi() - monkeypatch.setattr("pyharp.spectra.hitran_lines._import_hapi", lambda: fake) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils._import_hapi", lambda: fake) config = SpectroscopyConfig(output_path=tmp_path / "out.nc", hitran_cache_dir=tmp_path / "cache") band = SpectralBandConfig("single_state", 25.0, 2500.0, 1.0) config.hitran_cache_dir.mkdir(parents=True, exist_ok=True) @@ -106,7 +106,7 @@ def test_download_hitran_lines_refetches_contaminated_cache(monkeypatch, tmp_pat def test_download_hitran_lines_wraps_fetch_failures(monkeypatch, tmp_path) -> None: fake = FakeHapi() - monkeypatch.setattr("pyharp.spectra.hitran_lines._import_hapi", lambda: fake) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils._import_hapi", lambda: fake) def fail_fetch(table_name, iso_ids, numin, numax): raise RuntimeError("network down") @@ -134,9 +134,9 @@ def test_load_hitran_line_list_filters_by_min_line_strength(monkeypatch, tmp_pat "sw": [1.0e-30, 1.0e-27, 1.0e-25], } } - monkeypatch.setattr("pyharp.spectra.hitran_lines._import_hapi", lambda: fake) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils._import_hapi", lambda: fake) monkeypatch.setattr( - "pyharp.spectra.hitran_lines.download_hitran_lines", + "pyharp.spectra.hitran_molecule_utils.download_hitran_lines", lambda config, band: type("LineDb", (), {"table_name": table_name, "cache_dir": config.hitran_cache_dir})(), ) config = SpectroscopyConfig(output_path=tmp_path / "out.nc", hitran_cache_dir=tmp_path / "cache") @@ -150,7 +150,7 @@ def test_load_hitran_line_list_filters_by_min_line_strength(monkeypatch, tmp_pat def test_hapi_line_provider_passes_min_line_strength_to_hapi(monkeypatch, tmp_path) -> None: fake = FakeHapi() - monkeypatch.setattr("pyharp.spectra.hitran_lines._import_hapi", lambda: fake) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils._import_hapi", lambda: fake) fake.LOCAL_TABLE_CACHE["mock"] = {"data": {"gamma_self": [0.1]}} provider = HapiLineProvider("mock", cache_dir=tmp_path, min_line_strength=1.0e-27) @@ -163,7 +163,7 @@ def test_hapi_line_provider_passes_min_line_strength_to_hapi(monkeypatch, tmp_pa def test_hapi_line_provider_falls_back_missing_broadener_to_air(monkeypatch, tmp_path) -> None: fake = FakeHapi() fake.LOCAL_TABLE_CACHE["mock"] = {"data": {"gamma_air": [0.1], "gamma_self": [0.2]}} - monkeypatch.setattr("pyharp.spectra.hitran_lines._import_hapi", lambda: fake) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils._import_hapi", lambda: fake) provider = HapiLineProvider("mock", cache_dir=tmp_path, diluent={"h2": 0.7, "self": 0.3}) @@ -174,7 +174,7 @@ def test_hapi_line_provider_falls_back_missing_broadener_to_air(monkeypatch, tmp def test_hapi_line_provider_raises_when_air_fallback_unavailable(monkeypatch, tmp_path) -> None: fake = FakeHapi() fake.LOCAL_TABLE_CACHE["mock"] = {"data": {"gamma_self": [0.2]}} - monkeypatch.setattr("pyharp.spectra.hitran_lines._import_hapi", lambda: fake) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils._import_hapi", lambda: fake) try: HapiLineProvider("mock", cache_dir=tmp_path, diluent={"h2": 1.0}) @@ -187,7 +187,7 @@ def test_hapi_line_provider_raises_when_air_fallback_unavailable(monkeypatch, tm def test_build_line_provider_uses_config_broadening_composition(monkeypatch, tmp_path) -> None: fake = FakeHapi() fake.LOCAL_TABLE_CACHE["mock"] = {"data": {"gamma_air": [0.1], "gamma_self": [0.2], "gamma_h2": [0.3]}} - monkeypatch.setattr("pyharp.spectra.hitran_lines._import_hapi", lambda: fake) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils._import_hapi", lambda: fake) config = SpectroscopyConfig( output_path=tmp_path / "out.nc", diff --git a/tests/spectra/test_molecule_plot_cli.py b/tests/spectra/test_molecule_plot_cli.py index a602365..15ec99f 100644 --- a/tests/spectra/test_molecule_plot_cli.py +++ b/tests/spectra/test_molecule_plot_cli.py @@ -2,9 +2,8 @@ import numpy as np -from pyharp.spectra.hitran_lines import LineDatabase -from pyharp.spectra.molecule_plot_cli import _compute_overview_products, _compute_requested_absorption_spectrum -from pyharp.spectra.molecule_plot_cli import build_molecule_overview_batch_parser, build_overview_parser, run_xsection +from pyharp.spectra.hitran_molecule_plot import build_molecule_overview_batch_parser, build_overview_parser, compute_overview_products, run_xsection +from pyharp.spectra.hitran_molecule_utils import LineDatabase, compute_requested_absorption_spectrum def test_overview_parser_accepts_optional_cia_and_pdf_output(tmp_path) -> None: @@ -48,7 +47,7 @@ def test_batch_overview_parser_accepts_repeated_ranges_and_default_species(tmp_p args = parser.parse_args( [ "--wn-range=25,2500", - "--wn-range=2501,20000", + "--wn-range=2500,20000", "--temperature-k", "300", "--pressure-bar", @@ -63,7 +62,7 @@ def test_batch_overview_parser_accepts_repeated_ranges_and_default_species(tmp_p ) assert args.species == ["H2", "CO2", "H2O", "CH4", "N2"] - assert args.wn_ranges == [(25.0, 2500.0), (2501.0, 20000.0)] + assert args.wn_ranges == [(25.0, 2500.0), (2500.0, 20000.0)] assert args.temperature_k == 300.0 assert args.pressure_bar == 1.0 assert args.broadening_composition == "H2:0.85,He:0.15" @@ -149,17 +148,17 @@ def fake_compute_from_sources( pressure_pa=pressure_pa, ) - monkeypatch.setattr("pyharp.spectra.molecule_plot_cli.download_hitran_lines", fake_download) - monkeypatch.setattr("pyharp.spectra.molecule_plot_cli.load_hitran_line_list", fake_load_line_list) - monkeypatch.setattr("pyharp.spectra.molecule_plot_cli._load_requested_cia_dataset", lambda args, config: None) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils.download_hitran_lines", fake_download) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils.load_hitran_line_list", fake_load_line_list) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils.load_requested_cia_dataset", lambda args, config: None) monkeypatch.setattr( - "pyharp.spectra.molecule_plot_cli.build_line_provider", + "pyharp.spectra.hitran_molecule_utils.build_line_provider", lambda config, line_db: SimpleNamespace(broadening_summary=lambda: "requested=self:1.000 -> effective=self:1.000"), ) - monkeypatch.setattr("pyharp.spectra.molecule_plot_cli._resolve_continuum_sources", lambda **kwargs: (None, None)) - monkeypatch.setattr("pyharp.spectra.molecule_plot_cli.compute_absorption_spectrum_from_sources", fake_compute_from_sources) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils._resolve_continuum_sources", lambda **kwargs: (None, None)) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils.compute_absorption_spectrum_from_sources", fake_compute_from_sources) - _compute_overview_products(args) + compute_overview_products(args) assert calls["download"] == 1 @@ -194,14 +193,14 @@ def test_compute_requested_absorption_spectrum_reuses_built_line_provider_withou fake_provider = SimpleNamespace(broadening_summary=lambda: "requested=self:1.000 -> effective=self:1.000") calls = {"resolve": 0, "compute": 0} - monkeypatch.setattr("pyharp.spectra.molecule_plot_cli.download_hitran_lines", lambda config, band: line_db) - monkeypatch.setattr("pyharp.spectra.molecule_plot_cli._load_requested_cia_dataset", lambda args, config: None) - monkeypatch.setattr("pyharp.spectra.molecule_plot_cli.build_line_provider", lambda config, line_db: fake_provider) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils.download_hitran_lines", lambda config, band: line_db) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils.load_requested_cia_dataset", lambda args, config: None) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils.build_line_provider", lambda config, line_db: fake_provider) def fake_resolve_continuum_sources(*, config, wavenumber_grid_cm1, temperature_k, pressure_pa): calls["resolve"] += 1 assert config.hitran_species.name == "CO2" - assert np.allclose(wavenumber_grid_cm1, np.array([20.0, 21.0])) + assert np.allclose(wavenumber_grid_cm1, np.array([20.0, 21.0, 22.0])) assert temperature_k == 300.0 assert pressure_pa == 1.0e5 return None, None @@ -218,7 +217,7 @@ def fake_compute_from_sources( ): calls["compute"] += 1 assert species_name == "CO2" - assert np.allclose(wavenumber_grid_cm1, np.array([20.0, 21.0])) + assert np.allclose(wavenumber_grid_cm1, np.array([20.0, 21.0, 22.0])) assert temperature_k == 300.0 assert pressure_pa == 1.0e5 assert line_provider is fake_provider @@ -234,10 +233,10 @@ def fake_compute_from_sources( pressure_pa=pressure_pa, ) - monkeypatch.setattr("pyharp.spectra.molecule_plot_cli._resolve_continuum_sources", fake_resolve_continuum_sources) - monkeypatch.setattr("pyharp.spectra.molecule_plot_cli.compute_absorption_spectrum_from_sources", fake_compute_from_sources) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils._resolve_continuum_sources", fake_resolve_continuum_sources) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_utils.compute_absorption_spectrum_from_sources", fake_compute_from_sources) - _, _, _, line_provider = _compute_requested_absorption_spectrum(args) + _, _, _, line_provider = compute_requested_absorption_spectrum(args) assert calls == {"resolve": 1, "compute": 1} assert line_provider is fake_provider @@ -269,10 +268,10 @@ def test_run_xsection_reports_broadening_summary(monkeypatch, tmp_path, capsys) ) monkeypatch.setattr( - "pyharp.spectra.molecule_plot_cli._compute_requested_absorption_spectrum", + "pyharp.spectra.hitran_molecule_utils.compute_requested_absorption_spectrum", lambda *call_args, **call_kwargs: (None, None, fake_spectrum, fake_provider), ) - monkeypatch.setattr("pyharp.spectra.molecule_plot_cli.plot_absorption_spectrum", lambda spectrum, figure: None) + monkeypatch.setattr("pyharp.spectra.hitran_molecule_plot.plot_absorption_spectrum", lambda spectrum, figure: None) run_xsection(args) diff --git a/tests/spectra/test_output_names.py b/tests/spectra/test_output_names.py index 18760df..5257bdf 100644 --- a/tests/spectra/test_output_names.py +++ b/tests/spectra/test_output_names.py @@ -1,6 +1,6 @@ from pathlib import Path -from pyharp.spectra.output_names import default_output_path +from pyharp.spectra.utils import default_output_path def test_default_output_path_uses_requested_pattern() -> None: diff --git a/tests/spectra/test_plot_cli.py b/tests/spectra/test_plot_cli.py index 6797f7a..00ee797 100644 --- a/tests/spectra/test_plot_cli.py +++ b/tests/spectra/test_plot_cli.py @@ -87,7 +87,7 @@ def test_plot_parser_accepts_atm_overview_ranges(tmp_path) -> None: "--broadening-composition", "H2:0.85,He:0.15", "--wn-range=25,2500", - "--wn-range=2501,20000", + "--wn-range=2500,20000", "--manifest", str(tmp_path / "sources.json"), "--output", @@ -98,7 +98,7 @@ def test_plot_parser_accepts_atm_overview_ranges(tmp_path) -> None: assert args.command == "overview" assert args.composition == "H2O:0.1,H2:0.9" assert args.broadening_composition == "H2:0.85,He:0.15" - assert args.wn_ranges == [(25.0, 2500.0), (2501.0, 20000.0)] + assert args.wn_ranges == [(25.0, 2500.0), (2500.0, 20000.0)] assert args.manifest == tmp_path / "sources.json" assert args.figure == tmp_path / "atm.pdf" @@ -109,7 +109,7 @@ def test_plot_main_dispatches_pair_attenuation(monkeypatch) -> None: def fake_run(args): calls.append(args) - monkeypatch.setattr("pyharp.spectra.plot_cli.cia_plot_cli.run_attenuation", fake_run) + monkeypatch.setattr("pyharp.spectra.plot_cli.hitran_cia_plot.run_attenuation", fake_run) plot_cli.main(["attenuation", "--pair", "H2-He", "--wn-range", "25,30"]) @@ -125,7 +125,7 @@ def test_plot_main_dispatches_molecule_xsection_with_default_figure(monkeypatch) def fake_run(args): calls.append(args) - monkeypatch.setattr("pyharp.spectra.plot_cli.molecule_plot_cli.run_xsection", fake_run) + monkeypatch.setattr("pyharp.spectra.plot_cli.hitran_molecule_plot.run_xsection", fake_run) plot_cli.main(["xsection", "--species", "CO2", "--temperature-k", "275.5", "--pressure-bar", "0.25", "--wn-range", "25,30.5"]) @@ -139,7 +139,7 @@ def test_plot_main_dispatches_one_xsection_per_state_pair(monkeypatch, tmp_path) def fake_run(args): calls.append(args) - monkeypatch.setattr("pyharp.spectra.plot_cli.molecule_plot_cli.run_xsection", fake_run) + monkeypatch.setattr("pyharp.spectra.plot_cli.hitran_molecule_plot.run_xsection", fake_run) monkeypatch.setattr( "pyharp.spectra.plot_cli._parallel_plot_results", lambda tasks, *, worker: [worker(task) for task in tasks], @@ -173,7 +173,7 @@ def test_plot_main_appends_state_suffix_to_explicit_output_for_multiple_pairs(mo def fake_run(args): calls.append(args) - monkeypatch.setattr("pyharp.spectra.plot_cli.molecule_plot_cli.run_xsection", fake_run) + monkeypatch.setattr("pyharp.spectra.plot_cli.hitran_molecule_plot.run_xsection", fake_run) monkeypatch.setattr( "pyharp.spectra.plot_cli._parallel_plot_results", lambda tasks, *, worker: [worker(task) for task in tasks], @@ -205,7 +205,7 @@ def test_plot_main_uses_output_dir_for_default_figure(monkeypatch, tmp_path) -> def fake_run(args): calls.append(args) - monkeypatch.setattr("pyharp.spectra.plot_cli.molecule_plot_cli.run_xsection", fake_run) + monkeypatch.setattr("pyharp.spectra.plot_cli.hitran_molecule_plot.run_xsection", fake_run) plot_cli.main( [ @@ -233,7 +233,7 @@ def test_plot_main_passes_broadening_composition_to_molecule_workflow(monkeypatc def fake_run(args): calls.append(args) - monkeypatch.setattr("pyharp.spectra.plot_cli.molecule_plot_cli.run_xsection", fake_run) + monkeypatch.setattr("pyharp.spectra.plot_cli.hitran_molecule_plot.run_xsection", fake_run) plot_cli.main(["xsection", "--species", "CO2", "--broadening-composition", "air:0.8,self:0.2"]) @@ -266,13 +266,30 @@ def map(self, worker, tasks): assert created == {"max_workers": 2, "mp_context": "ctx-token"} +def test_parallel_plot_results_falls_back_to_subprocess_workers_when_process_pool_locks_are_unavailable(monkeypatch) -> None: + class FailingExecutor: + def __init__(self, *args, **kwargs): + raise PermissionError(13, "Permission denied") + + monkeypatch.setattr("pyharp.spectra.plot_cli.process_pool_context", lambda: "ctx-token") + monkeypatch.setattr("pyharp.spectra.plot_cli.ProcessPoolExecutor", FailingExecutor) + monkeypatch.setattr( + "pyharp.spectra.plot_cli._parallel_plot_results_via_subprocess", + lambda tasks, *, worker, max_workers: ["subprocess-result"], + ) + + result = plot_cli._parallel_plot_results([("x", 1), ("y", 2)], worker=plot_cli._run_plot_task) + + assert result == ["subprocess-result"] + + def test_plot_main_dispatches_composition_attenuation(monkeypatch) -> None: calls = [] def fake_run(args, *, wn_range): calls.append((args, wn_range)) - monkeypatch.setattr("pyharp.spectra.plot_cli.atm_overview_cli.run_atm_attenuation", fake_run) + monkeypatch.setattr("pyharp.spectra.plot_cli.atm_overview.run_atm_attenuation", fake_run) plot_cli.main(["attenuation", "--composition", "H2:0.9,He:0.1", "--wn-range", "25,30"]) @@ -288,7 +305,7 @@ def test_plot_main_preserves_explicit_figure_over_output_dir(monkeypatch, tmp_pa def fake_run(args): calls.append(args) - monkeypatch.setattr("pyharp.spectra.plot_cli.cia_plot_cli.run_binary", fake_run) + monkeypatch.setattr("pyharp.spectra.plot_cli.hitran_cia_plot.run_binary", fake_run) plot_cli.main( [ @@ -360,7 +377,7 @@ def test_plot_composition_transmission_matches_dump_total(monkeypatch, tmp_path) def fake_run(args, *, wn_range): calls.append((args, wn_range)) - monkeypatch.setattr("pyharp.spectra.plot_cli.atm_overview_cli.run_atm_transmission", fake_run) + monkeypatch.setattr("pyharp.spectra.plot_cli.atm_overview.run_atm_transmission", fake_run) composition = "H2:0.9,He:0.1,H2O:0.002" plot_cli.main( @@ -441,12 +458,12 @@ def savefig(self, fig): dummy_figure = type("Figure", (), {})() dummy_axis = type("Axis", (), {})() - monkeypatch.setattr("pyharp.spectra.atm_overview_cli._parallel_mixture_overview_products", fake_parallel_products) - monkeypatch.setattr("pyharp.spectra.atm_overview_cli._render_mixture_overview", lambda fig, axes, *, products: None) - monkeypatch.setattr("pyharp.spectra.atm_overview_cli._page_manifest", lambda products: {}) - monkeypatch.setattr("pyharp.spectra.atm_overview_cli.PdfPages", _DummyPdf) - monkeypatch.setattr("pyharp.spectra.atm_overview_cli.plt.subplots", lambda **kwargs: (dummy_figure, np.array([[dummy_axis], [dummy_axis], [dummy_axis], [dummy_axis]]))) - monkeypatch.setattr("pyharp.spectra.atm_overview_cli.plt.close", lambda fig: None) + monkeypatch.setattr("pyharp.spectra.atm_overview._parallel_mixture_overview_products", fake_parallel_products) + monkeypatch.setattr("pyharp.spectra.atm_overview._render_mixture_overview", lambda fig, axes, *, products: None) + monkeypatch.setattr("pyharp.spectra.atm_overview._page_manifest", lambda products: {}) + monkeypatch.setattr("pyharp.spectra.atm_overview.PdfPages", _DummyPdf) + monkeypatch.setattr("pyharp.spectra.atm_overview.plt.subplots", lambda **kwargs: (dummy_figure, np.array([[dummy_axis], [dummy_axis], [dummy_axis], [dummy_axis]]))) + monkeypatch.setattr("pyharp.spectra.atm_overview.plt.close", lambda fig: None) plot_cli.main( [ @@ -474,7 +491,7 @@ def test_plot_overview_uses_single_output_pdf_for_multiple_state_pairs(monkeypat def fake_run(args): calls.append(args) - monkeypatch.setattr("pyharp.spectra.plot_cli.molecule_plot_cli.run_overview", fake_run) + monkeypatch.setattr("pyharp.spectra.plot_cli.hitran_molecule_plot.run_overview", fake_run) plot_cli.main( [ @@ -502,7 +519,7 @@ def test_plot_overview_uses_combined_default_name_without_duplicate_units(monkey def fake_run(args): calls.append(args) - monkeypatch.setattr("pyharp.spectra.plot_cli.molecule_plot_cli.run_overview_batch", fake_run) + monkeypatch.setattr("pyharp.spectra.plot_cli.hitran_molecule_plot.run_overview_batch", fake_run) plot_cli.main( [ diff --git a/tests/spectra/test_shared_cli.py b/tests/spectra/test_shared_cli.py index 61068f8..f8ad6c7 100644 --- a/tests/spectra/test_shared_cli.py +++ b/tests/spectra/test_shared_cli.py @@ -1,36 +1,36 @@ import types -from pyharp.spectra import shared_cli +from pyharp.spectra import utils def test_process_pool_context_uses_spawn_on_macos(monkeypatch) -> None: - monkeypatch.setattr(shared_cli.sys, "platform", "darwin") + monkeypatch.setattr(utils.sys, "platform", "darwin") calls = [] def fake_get_context(method): calls.append(method) return types.SimpleNamespace(method=method) - monkeypatch.setattr(shared_cli.mp, "get_context", fake_get_context) + monkeypatch.setattr(utils.mp, "get_context", fake_get_context) - ctx = shared_cli.process_pool_context() + ctx = utils.process_pool_context() assert ctx.method == "spawn" assert calls == ["spawn"] def test_process_pool_context_prefers_fork_when_available(monkeypatch) -> None: - monkeypatch.setattr(shared_cli.sys, "platform", "linux") - monkeypatch.setattr(shared_cli.mp, "get_all_start_methods", lambda: ["fork", "spawn"]) + monkeypatch.setattr(utils.sys, "platform", "linux") + monkeypatch.setattr(utils.mp, "get_all_start_methods", lambda: ["fork", "spawn"]) calls = [] def fake_get_context(method): calls.append(method) return types.SimpleNamespace(method=method) - monkeypatch.setattr(shared_cli.mp, "get_context", fake_get_context) + monkeypatch.setattr(utils.mp, "get_context", fake_get_context) - ctx = shared_cli.process_pool_context() + ctx = utils.process_pool_context() assert ctx.method == "fork" assert calls == ["fork"] diff --git a/tests/spectra/test_spectrum.py b/tests/spectra/test_spectrum.py index ff99ea4..f12e037 100644 --- a/tests/spectra/test_spectrum.py +++ b/tests/spectra/test_spectrum.py @@ -3,7 +3,7 @@ import numpy as np import xarray as xr -from pyharp.spectra.hitran_cia import CiaBlock, CiaDataset +from pyharp.spectra.hitran_cia_utils import CiaBlock, CiaDataset from pyharp.spectra.spectrum import ( compute_absorption_spectrum_from_sources, plot_absorption_spectrum, diff --git a/tests/test_attenuator.cpp b/tests/test_attenuator.cpp index 20b0ba3..9cfbaff 100644 --- a/tests/test_attenuator.cpp +++ b/tests/test_attenuator.cpp @@ -1,111 +1,214 @@ +// C/C++ +#include +#include +#include + +// base +#include + // external #include // harp -#include -#include -#include -#include -#include - -using namespace harp; - -TEST(TestOpacity, sw) { - AttenuatorOptions op; - species_names = {"S8", "H2SO4"}; - species_weights = {256.0e-3, 98.0e-3}; - - op.species_ids({0}).opacity_files({"s8_k_fuller.txt"}); - S8Fuller s8(op); - - op.species_ids({1}).opacity_files({"h2so4.txt"}); - H2SO4Simple h2so4(op); - - // std::cout << "h2so4 wave = " << h2so4->kwave << std::endl; - // std::cout << "h2so4 data = " << h2so4->kdata << std::endl; - - int ncol = 1; - int nlyr = 1; - int nspecies = 2; - auto conc = torch::ones({ncol, nlyr, nspecies}, torch::kFloat64); - std::map kwargs; - kwargs["wavelength"] = s8->kwave; - - auto result1 = s8->forward(conc, kwargs); - auto result2 = h2so4->forward(conc, kwargs); - // std::cout << "result1 = " << result1 << std::endl; - // std::cout << "result2 = " << result2 << std::endl; - // std::cout << result2.sizes() << std::endl; - - // attenuation [1/m] - auto result = result1 + result2; - - // single scattering albedo - result.select(-1, 1) = result.select(-1, 1) / result.select(-1, 0); +#include - std::cout << "result = " << result << std::endl; -}; +#include +#include +#include -TEST(TestOpacity, lw) { - AttenuatorOptions op; - species_names = {"CO2", "H2O"}; - species_weights = {44.0e-3, 18.0e-3}; - - op.species_ids({0}).opacity_files({"amarsw-ck-B1.nc"}); - RFM co2(op); - - std::cout << "co2 = " << co2->kdata << std::endl; - std::cout << "pres = " << exp(co2->klnp) << std::endl; - std::cout << "temp = " << co2->ktempa << std::endl; - std::cout << "wave = " << co2->kwave << std::endl; - - op.species_ids({1}).opacity_files({"amarsw-ck-B1.nc"}); - RFM h2o(op); +// netcdf +#ifdef NETCDFOUTPUT +extern "C" { +#include } +#endif -TEST(TestOpacity, get_reftemp) { - AttenuatorOptions op; - species_names = {"CO2", "H2O"}; - species_weights = {44.0e-3, 18.0e-3}; +namespace fs = std::filesystem; - op.species_ids({0}).opacity_files({"amarsw-ck-B1.nc"}); - RFM co2(op); +namespace harp { +extern std::vector species_names; +extern std::vector species_weights; +} // namespace harp - auto coord = (torch::ones({2, 2}, torch::kFloat64) * 100.e5).log(); - std::cout << "coord = " << coord << std::endl; +namespace { - auto temp = interpn({coord}, {co2->klnp}, co2->kreftem).squeeze(-1); - std::cout << "temp = " << temp << std::endl; +#ifdef NETCDFOUTPUT +void check_nc(int status) { + ASSERT_EQ(status, NC_NOERR) << nc_strerror(status); } -TEST(TestOpacity, forward) { - AttenuatorOptions op; - species_names = {"CO2", "H2O"}; - species_weights = {44.0e-3, 18.0e-3}; - - op.species_ids({0}).opacity_files({"amarsw-ck-B1.nc"}); - RFM co2(op); - - op.species_ids({1}).opacity_files({"amarsw-ck-B1.nc"}); - RFM h2o(op); +void put_text_attr(int fileid, int varid, char const* name, char const* value) { + check_nc(nc_put_att_text(fileid, varid, name, std::strlen(value), value)); +} - int ncol = 1; - int nlyr = 1; - int nspecies = 2; - auto conc = torch::ones({ncol, nlyr, nspecies}, torch::kFloat64); - std::map kwargs; - kwargs["pres"] = torch::ones({ncol, nlyr}, torch::kFloat64) * 10.e5; - kwargs["temp"] = torch::ones({ncol, nlyr}, torch::kFloat64) * 300.0; +fs::path write_test_dataset() { + auto path = fs::temp_directory_path() / "pyharp_test_molecule_line.nc"; + int fileid = -1; + check_nc(nc_create(path.c_str(), NC_CLOBBER, &fileid)); + + int dim_del_temp = -1, dim_pressure = -1, dim_wavenumber = -1; + check_nc(nc_def_dim(fileid, "del_temperature", 2, &dim_del_temp)); + check_nc(nc_def_dim(fileid, "pressure", 2, &dim_pressure)); + check_nc(nc_def_dim(fileid, "wavenumber", 3, &dim_wavenumber)); + + int var_wavenumber = -1, var_pressure = -1, var_del_temp = -1; + int var_temperature = -1, var_line = -1, var_cont = -1, var_cia = -1; + check_nc(nc_def_var(fileid, "wavenumber", NC_DOUBLE, 1, &dim_wavenumber, + &var_wavenumber)); + check_nc(nc_def_var(fileid, "pressure", NC_DOUBLE, 1, &dim_pressure, + &var_pressure)); + check_nc(nc_def_var(fileid, "del_temperature", NC_DOUBLE, 1, &dim_del_temp, + &var_del_temp)); + check_nc(nc_def_var(fileid, "temperature", NC_DOUBLE, 1, &dim_pressure, + &var_temperature)); + + int dims3[3] = {dim_del_temp, dim_pressure, dim_wavenumber}; + check_nc( + nc_def_var(fileid, "sigma_line_h2o", NC_DOUBLE, 3, dims3, &var_line)); + check_nc(nc_def_var(fileid, "sigma_continuum_h2o_mt_ckd", NC_DOUBLE, 3, dims3, + &var_cont)); + check_nc(nc_def_var(fileid, "binary_absorption_coefficient_h2_he", NC_DOUBLE, + 3, dims3, &var_cia)); + + put_text_attr(fileid, var_wavenumber, "units", "cm^-1"); + put_text_attr(fileid, var_pressure, "units", "Pa"); + put_text_attr(fileid, var_del_temp, "units", "K"); + put_text_attr(fileid, var_temperature, "units", "K"); + put_text_attr(fileid, var_line, "units", "cm^2 molecule^-1"); + put_text_attr(fileid, var_cont, "units", "cm^2 molecule^-1"); + put_text_attr(fileid, var_cia, "units", "cm^5 molecule^-2"); + + check_nc(nc_enddef(fileid)); + + double const wavenumber[] = {20.0, 21.0, 22.0}; + double const pressure[] = {1.0e5, 1.0e6}; + double const del_temp[] = {-10.0, 10.0}; + double const temperature[] = {300.0, 500.0}; + + std::vector sigma_line(2 * 2 * 3); + std::vector sigma_cont(2 * 2 * 3); + std::vector sigma_cia(2 * 2 * 3); + + for (int idt = 0; idt < 2; ++idt) { + for (int ip = 0; ip < 2; ++ip) { + for (int iw = 0; iw < 3; ++iw) { + auto idx = (idt * 2 + ip) * 3 + iw; + sigma_line[idx] = (1.0 + idx) * 1.0e-24; + sigma_cont[idx] = (0.1 + idx) * 1.0e-24; + sigma_cia[idx] = (2.0 + idx) * 1.0e-46; + } + } + } + + sigma_line[0] = 0.0; + sigma_cont[0] = 0.0; + sigma_cia[0] = 0.0; + + check_nc(nc_put_var_double(fileid, var_wavenumber, wavenumber)); + check_nc(nc_put_var_double(fileid, var_pressure, pressure)); + check_nc(nc_put_var_double(fileid, var_del_temp, del_temp)); + check_nc(nc_put_var_double(fileid, var_temperature, temperature)); + check_nc(nc_put_var_double(fileid, var_line, sigma_line.data())); + check_nc(nc_put_var_double(fileid, var_cont, sigma_cont.data())); + check_nc(nc_put_var_double(fileid, var_cia, sigma_cia.data())); + check_nc(nc_close(fileid)); + + return path; +} +#endif + +TEST(TestOpacity, MoleculeLineAddsContinuumAndHandlesDimensionOrder) { +#ifndef NETCDFOUTPUT + GTEST_SKIP() << "NetCDF support is disabled"; +#else + auto dataset = write_test_dataset(); + harp::species_names = {"H2O", "H2", "He"}; + harp::species_weights = {18.0e-3, 2.0e-3, 4.0e-3}; + + auto op = harp::OpacityOptionsImpl::create(); + op->type("molecule-line").species_ids({0}).opacity_files({dataset.string()}); + harp::MoleculeLine line(op); + + auto conc = torch::zeros({1, 1, 3}, torch::kFloat64); + conc[0][0][0] = 2.0; + std::map atm; + atm["pres"] = torch::tensor({{1.0e5}}, torch::kFloat64); + atm["temp"] = torch::tensor({{290.0}}, torch::kFloat64); + atm["wavenumber"] = torch::tensor({20.0, 21.0, 22.0}, torch::kFloat64); + + auto result = line->forward(conc, atm).squeeze(-1).squeeze(-1).squeeze(-1); + auto expected_sigma = + torch::tensor({3.1e-24, 3.1e-24, 5.1e-24}, torch::kFloat64) * + (1.0e-4 * harp::constants::Avogadro); + auto expected = expected_sigma * 2.0; + EXPECT_TRUE(torch::allclose(result, expected, 1.0e-12, 1.0e-12)); +#endif +} - auto result1 = co2->forward(conc, kwargs); - auto result2 = h2o->forward(conc, kwargs); +TEST(TestOpacity, CIAHandlesBinaryPairsAndReversedSpeciesOrder) { +#ifndef NETCDFOUTPUT + GTEST_SKIP() << "NetCDF support is disabled"; +#else + auto dataset = write_test_dataset(); + harp::species_names = {"H2O", "H2", "He"}; + harp::species_weights = {18.0e-3, 2.0e-3, 4.0e-3}; + + auto op = harp::OpacityOptionsImpl::create(); + op->type("molecule-cia") + .species_ids({2, 1}) + .opacity_files({dataset.string()}); + harp::MoleculeCIA cia(op); + + auto conc = torch::zeros({1, 1, 3}, torch::kFloat64); + conc[0][0][1] = 3.0; + conc[0][0][2] = 4.0; + std::map atm; + atm["pres"] = torch::tensor({{1.0e5}}, torch::kFloat64); + atm["temp"] = torch::tensor({{290.0}}, torch::kFloat64); + atm["wavenumber"] = torch::tensor({20.0, 21.0, 22.0}, torch::kFloat64); + + auto result = cia->forward(conc, atm).squeeze(-1).squeeze(-1).squeeze(-1); + auto expected_coeff = + torch::tensor({3.0e-46, 3.0e-46, 4.0e-46}, torch::kFloat64) * + (1.0e-10 * harp::constants::Avogadro * harp::constants::Avogadro); + auto expected = expected_coeff * 12.0; + EXPECT_TRUE(torch::allclose(result, expected, 1.0e-12, 1.0e-12)); +#endif +} - std::cout << "result1 = " << result1 << std::endl; - std::cout << "result2 = " << result2.squeeze(1).squeeze(1) << std::endl; +TEST(TestOpacity, NewOpacityTypesParseFromYaml) { + harp::species_names = {"H2O", "H2", "He"}; + harp::species_weights = {18.0e-3, 2.0e-3, 4.0e-3}; + + auto yaml_path = fs::temp_directory_path() / "pyharp_test_opacity.yaml"; + std::ofstream out(yaml_path); + out << "opacities:\n" + << " line:\n" + << " type: molecule-line\n" + << " data: [/tmp/mock.nc]\n" + << " species: [H2O]\n" + << " cia_pair:\n" + << " type: molecule-cia\n" + << " data: [/tmp/mock.nc]\n" + << " species: [H2, He]\n"; + out.close(); + + auto line = harp::OpacityOptionsImpl::from_yaml(yaml_path.string(), "line"); + EXPECT_EQ(line->type(), "molecule-line"); + ASSERT_EQ(line->species_ids().size(), 1); + EXPECT_EQ(line->species_ids()[0], 0); + + auto cia = + harp::OpacityOptionsImpl::from_yaml(yaml_path.string(), "cia_pair"); + EXPECT_EQ(cia->type(), "molecule-cia"); + ASSERT_EQ(cia->species_ids().size(), 2); + EXPECT_EQ(cia->species_ids()[0], 1); + EXPECT_EQ(cia->species_ids()[1], 2); } +} // namespace + int main(int argc, char** argv) { testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); }