more docs and refractoring

.flake8

@@ -10,4 +10,5 @@ exclude =
     old,
     build,
     dist,
-    tests/*.py
+    tests/*.py
+    opac_mixer/patches/*.py

README.md

@@ -7,9 +7,4 @@ The documentation for the code and tutorials are provided here: [link to docs](h
 If you find this work useful, please consider to cite the following paper:
 Schneider et al. (in review)
 
-WIP:
-- [ ] extend the docs
-- [ ] polish up the code and add more docstrings to the code
-- [ ] publish python package on pypi
-
 WARNING: The paper is not yet published, please do not use this code, before it is published

doc/source/API.rst

@@ -0,0 +1,33 @@
+.. prt_phasecurve documentation master file, created by
+   sphinx-quickstart on Mon Jan 11 14:18:49 2021.
+   You can adapt this file completely to your liking, but it should at least
+   contain the root `toctree` directive.
+
+API Documentation
+=================
+
+emulator.py
+-----------
+
+ .. automodule:: opac_mixer.emulator
+
+ .. autoclass:: Emulator
+   :members:
+
+read.py
+-------
+
+.. automodule:: opac_mixer.read
+
+ .. autoclass:: ReadOpac
+   :members:
+.. autoclass:: ReadOpacChubb
+   :members:
+
+mix.py
+------
+
+.. automodule:: opac_mixer.mix
+
+ .. autoclass:: CombineOpacGrid
+   :members:

doc/source/Coupling.rst

@@ -0,0 +1,66 @@
+.. prt_phasecurve documentation master file, created by
+   sphinx-quickstart on Mon Jan 11 14:18:49 2021.
+   You can adapt this file completely to your liking, but it should at least
+   contain the root `toctree` directive.
+
+Deployment in code
+==================
+
+Coupling the DeepSet mixing to a radiative transfer solver is easy to do, if you have already some mixing in your radiative transfer solver.
+It requires two steps:
+
+1. Reading in your custom k-tables (see :ref:`Tutorial: Add custom k-tables`) from your radiative transfer solver and training the DeepSet on these (see :ref:`Tutorial: Quick Start`).
+
+2. Implementing the inference of the neural network in your code
+
+Here, we will talk about the second step, since the first step is defined at other places in this documentation.
+
+Numpy implementation
+--------------------
+
+First, lets start with a simple numpy implementation of the DeepSet:
+
+.. code-block::
+   python
+
+   mlp_weights = [weights.numpy() for weights in em.model.weights]
+   def simple_mlp(kappas):
+       rep = np.tensordot(kappas, mlp_weights[0], axes=(1,0))  # first dense
+       rep[rep <= 0.0] = 0.0
+       sum_rep = np.sum(rep, axis=1)   # sum
+       dec = np.tensordot(sum_rep, mlp_weights[1], axes=(-1,0))  # second dense
+       return dec
+
+.. Note::
+
+   We use input and output scaling! It is therefore important to feed the scaled input to the function and to transform the output back using the inverse output scaling.
+
+General remarks
+---------------
+
+The exact deployment deployment depends on the code, language and framework you want to couple it to.
+Some general recommendations can be found in the paper (Appendix).
+
+Basically, the following things are always needed:
+
+0. (Implement the read in and interpolation of the individual k-tables and the chemistry in your model)
+
+1. Export the weights (there is an ``export`` function on the ``Emulator`` class)
+
+2. Read the weights into your model
+
+3. Implement the scaling functions (see ``opac_mixer/utils/scaling.py``)
+
+4. Implement the DeepSet mixing consisting of:
+
+   - A first matrix vector multiplication for each of the individual species
+
+   - A relu activation
+
+   - A sum over all hidden representations
+
+   - Another matrix vector multiplication
+
+.. Note::
+
+   It would generally be the fastest option to stack the matrix vector multiplications and deploy them using a vectorized matrix vector multiplication. Keep that in mind.

doc/source/Installation.rst

@@ -1,10 +1,16 @@
 Installation
 ------------
 
+Make sure you use a fresh anaconda environment, so that you don't mess up anything.
+
+The code has minimal requirements. The biggest requirement is ``tensorflow``, which should be installed in preparation.
+
+If you want to couple it to ``petitRADTRANS``, you might consider installing ``petitRADTRANS`` as well
+
 pip
 ^^^
 
-Future: once the project is on pypi, you can install it via
+You can install the code from pypi
 
 .. code-block:: bash
 

doc/source/conf.py

@@ -15,7 +15,6 @@
 
 sys.path.insert(0, os.path.abspath("../.."))
 
-
 # -- Project information -----------------------------------------------------
 
 project = "opac_mixer"
@@ -25,7 +24,6 @@
 # The full version, including alpha/beta/rc tags
 release = "0.1"
 
-
 # -- General configuration ---------------------------------------------------
 
 # Add any Sphinx extension module names here, as strings. They can be
@@ -68,6 +66,8 @@
     "numpy",
     "matplotlib",
     "tqdm",
+    "sklearn",
+    "keras",
     "h5py",
     "tensorflow",
     "MITgcmutils",
@@ -89,3 +89,13 @@
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
 html_static_path = ["_static"]
+
+
+def skip(app, what, name, obj, would_skip, options):
+    if name == "__init__":
+        return False
+    return would_skip
+
+
+def setup(app):
+    app.connect("autodoc-skip-member", skip)

doc/source/index.rst

@@ -5,23 +5,30 @@
 
 Welcome to opac_mixer's documentation!
 ==========================================
-opacity mixing accelerated
+There is a pressing need for fast opacity mixing in the correlated-k approximation.
+This python package builds a framework for machine learning on grids of k-tables.
 
+To get the most out of this code, start by installing the code (see :ref:`Installation`).
+Once installed, head over to :ref:`Tutorial: Quick Start` to see a quick explanation of how to use the code.
+
+Finally, if you want to couple the code to your own radiative transfer, you might need to read in your own k-tables.
+Thats easy, and its explained in :ref:`Tutorial: Add custom k-tables`. You also need to write a short code passage to deploy the mixing in the radiative transfer solver (see :ref:`Deployment in code`).
+
+If you find this work useful, please consider to cite the following paper:
+Schneider et al. (in review)
+
+WARNING: The paper is not yet published, please do not use this code, before it is published
 
 .. toctree::
    :maxdepth: 0
    :caption: Contents
 
    Installation
    notebooks/training.ipynb
+   Coupling
+   notebooks/extend_reader.ipynb
+   API
    notebooks/petitRADTRANS.ipynb
    notebooks/hp_tuning.ipynb
-   notebooks/extend_reader.ipynb
-
-Indices and tables
-==================
-
-* :ref:`genindex`
-* :ref:`search`
 
 Copyright 2023 Aaron Schneider. Feel free to contact me for any questions via `mail <mailto:[email protected]>`_.

doc/source/notebooks/hp_tuning.ipynb

@@ -37,6 +37,15 @@
     "from hyperopt import hp, fmin, tpe, Trials, STATUS_OK, STATUS_FAIL"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# Tutorial: Hyperparameter tuning"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
   {
    "cell_type": "code",
    "execution_count": 4,
@@ -46,9 +55,6 @@
      "start_time": "2023-05-02T11:18:26.406583Z"
     },
     "collapsed": false,
-    "jupyter": {
-     "outputs_hidden": false
-    },
     "tags": []
    },
    "outputs": [],