diff --git a/.readthedocs.yaml b/.readthedocs.yaml
index 72d4a52..53d25a9 100644
--- a/.readthedocs.yaml
+++ b/.readthedocs.yaml
@@ -18,6 +18,7 @@ build:
 # Build documentation in the "docs/" directory with Sphinx
 sphinx:
    configuration: docs/source/conf.py
+   builder: html
 
 # Optionally build your docs in additional formats such as PDF and ePub
 # formats:
@@ -27,6 +28,7 @@ sphinx:
 # Optional but recommended, declare the Python requirements required
 # to build your documentation
 # See https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html
-# python:
-#    install:
-#    - requirements: docs/requirements.txt
\ No newline at end of file
+# Optionally set the version of Python and requirements required to build your docs
+python:
+  install:
+    - requirements: requirements-docs.txt
\ No newline at end of file
diff --git a/docs/source/FAQ.rst b/docs/source/FAQ.rst
index 653123a..189c374 100644
--- a/docs/source/FAQ.rst
+++ b/docs/source/FAQ.rst
@@ -3,4 +3,4 @@ FAQ
 
 *Have questions?*
 
-Please reach out and ask, we will add the common ones here.
\ No newline at end of file
+Please file an issue at https://github.com/ergodicio/inverse-thomson-scattering/issues and we can chat there!
\ No newline at end of file
diff --git a/docs/source/TSfitter.rst b/docs/source/TSfitter.rst
new file mode 100644
index 0000000..78bc7be
--- /dev/null
+++ b/docs/source/TSfitter.rst
@@ -0,0 +1,8 @@
+TS Fitter API
+================
+
+This is the module that handles the loss function, gradient calculation, hessian calculation, and parameter generation
+
+.. autoclass:: inverse_thomson_scattering.model.TSFitter.TSFitter
+    :members:
+    :private-members:
diff --git a/docs/source/api_main.rst b/docs/source/api_main.rst
new file mode 100644
index 0000000..205bc2f
--- /dev/null
+++ b/docs/source/api_main.rst
@@ -0,0 +1,34 @@
+API Documentation
+===================
+
+Note: This section is under heavy development and the API is very much subject to change
+
+
+High Level API
+---------------
+.. toctree::
+    :maxdepth: 2
+    :caption: Contents:
+
+    runner
+    fitter
+
+Middle (?) Level API
+----------------------
+.. toctree::
+    :maxdepth: 2
+    :caption: Contents:
+
+    TSfitter
+    spectrum
+
+
+Low Level API
+--------------
+.. toctree::
+    :maxdepth: 2
+    :caption: Contents:
+
+    formfactor
+    fitmodel
+    ratintn
\ No newline at end of file
diff --git a/docs/source/conf.py b/docs/source/conf.py
index a3abb09..04620fd 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -5,23 +5,280 @@
 
 # -- Project information -----------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
+# -- Path setup --------------------------------------------------------------
+
+# If extensions (or modules to document with autodoc) are in another directory,
+# add these directories to sys.path here. If the directory is relative to the
+# documentation root, use os.path.abspath to make it absolute, like shown here.
+#
+# import os
+# import sys
+# sys.path.insert(0, os.path.abspath('.'))
+
+
+# -- Project information -----------------------------------------------------
+
+import inspect
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+sys.path.insert(0, os.path.abspath("."))
+sys.path.append(os.path.abspath("../../"))
+
+import inverse_thomson_scattering
 
 project = 'TSADAR'
 copyright = '2023, Avi Milder, Archis Joglekar'
 author = 'Avi Milder, Archis Joglekar'
 
+# -- Options for HTML output -------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
+
+html_theme = 'sphinx_rtd_theme'
+html_static_path = ['_static']
+
+
+
 # -- General configuration ---------------------------------------------------
-# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
 
-extensions = []
+# Add any Sphinx extension module names here, as strings. They can be
+# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
+# ones.
+extensions = [
+    "sphinx.ext.autodoc",
+    "sphinx.ext.autosummary",
+    "sphinx.ext.coverage",
+    "sphinx.ext.doctest",
+    "sphinx.ext.githubpages",
+    "sphinx.ext.intersphinx",
+    "sphinx.ext.mathjax",
+    "sphinx.ext.napoleon",
+    "sphinx_copybutton",
+    "sphinx_github_style",
+]
+# options for sphinx_github_style
+top_level = "inverse_thomson_scattering"
+linkcode_blob = "head"
+linkcode_url = r"https://github.com/ergodicio/inverse-thomson-scattering/"
+linkcode_link_text = "Source"
+
+
+def linkcode_resolve(domain, info):
+    """Returns a link to the source code on GitHub, with appropriate lines highlighted"""
+
+    if domain != "py" or not info["module"]:
+        return None
+
+    modname = info["module"]
+    fullname = info["fullname"]
+
+    submod = sys.modules.get(modname)
+    if submod is None:
+        return None
+
+    obj = submod
+    for part in fullname.split("."):
+        try:
+            obj = getattr(obj, part)
+        except AttributeError:
+            return None
+
+    # for jitted stuff, get the original src
+    if hasattr(obj, "__wrapped__"):
+        obj = obj.__wrapped__
+
+    # get the link to HEAD
+    cmd = "git log -n1 --pretty=%H"
+    try:
+        # get most recent commit hash
+        head = subprocess.check_output(cmd.split()).strip().decode("utf-8")
+
+        # if head is a tag, use tag as reference
+        cmd = "git describe --exact-match --tags " + head
+        try:
+            tag = subprocess.check_output(cmd.split(" ")).strip().decode("utf-8")
+            blob = tag
+
+        except subprocess.CalledProcessError:
+            blob = head
+
+    except subprocess.CalledProcessError:
+        print("Failed to get head")  # so no head?
+        blob = "main"
+
+    linkcode_url = r"https://github.com/ergodicio/inverse-thomson-scattering/"
+    linkcode_url = linkcode_url.strip("/") + f"/blob/{blob}/"
+    linkcode_url += "{filepath}#L{linestart}-L{linestop}"
+
+    # get a Path object representing the working directory of the repository.
+    try:
+        cmd = "git rev-parse --show-toplevel"
+        repo_dir = Path(subprocess.check_output(cmd.split(" ")).strip().decode("utf-8"))
+
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError("Unable to determine the repository directory") from e
+
+    # For ReadTheDocs, repo is cloned to /path/to/<repo_dir>/checkouts/<version>/
+    if repo_dir.parent.stem == "checkouts":
+        repo_dir = repo_dir.parent.parent
+
+    # path to source file
+    try:
+        filepath = os.path.relpath(inspect.getsourcefile(obj), repo_dir)
+        if filepath is None:
+            return
+    except Exception:
+        return None
 
-templates_path = ['_templates']
-exclude_patterns = []
+    # lines in source file
+    try:
+        source, lineno = inspect.getsourcelines(obj)
+    except OSError:
+        return None
+    else:
+        linestart, linestop = lineno, lineno + len(source) - 1
 
+    # Fix links with "../../../" or "..\\..\\..\\"
+    filepath = "/".join(filepath[filepath.find(top_level) :].split("\\"))
+
+    final_link = linkcode_url.format(
+        filepath=filepath, linestart=linestart, linestop=linestop
+    )
+    print(f"Final Link for {fullname}: {final_link}")
+    return final_link
+
+
+# numpydoc_class_members_toctree = False
+# Napoleon settings
+napoleon_google_docstring = True
+napoleon_numpy_docstring = False
+napoleon_include_init_with_doc = False
+napoleon_include_private_with_doc = False
+napoleon_include_special_with_doc = True
+napoleon_use_admonition_for_examples = False
+napoleon_use_admonition_for_notes = False
+napoleon_use_admonition_for_references = False
+napoleon_use_ivar = False
+napoleon_use_param = True
+napoleon_use_rtype = False
+
+autodoc_default_options = {
+    "member-order": "bysource",
+    "special-members": "__call__",
+    "exclude-members": "__init__",
+}
+# Add any paths that contain templates here, relative to this directory.
+templates_path = ["_templates"]
+
+# The suffix(es) of source filenames.
+# You can specify multiple suffix as a list of string:
+#
+source_suffix = [".rst", ".md"]
+# source_suffix = {
+#     '.rst': 'restructuredtext',
+#     '.md': 'markdown',
+# }
+# The master toctree document.
+master_doc = "index"
+
+# The language for content autogenerated by Sphinx. Refer to documentation
+# for a list of supported languages.
+#
+# This is also used if you do content translation via gettext catalogs.
+# Usually you set "language" from the command line for these cases.
+language = "en"
+
+# List of patterns, relative to source directory, that match files and
+# directories to ignore when looking for source files.
+# This pattern also affects html_static_path and html_extra_path.
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store", "README.rst"]
 
 
 # -- Options for HTML output -------------------------------------------------
-# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
 
-html_theme = 'alabaster'
-html_static_path = ['_static']
+# The theme to use for HTML and HTML Help pages.  See the documentation for
+# a list of builtin themes.
+#
+# html_theme = "alabaster"
+
+# Theme options are theme-specific and customize the look and feel of a theme
+# further.  For a list of options available for each theme, see the
+# documentation.
+
+html_theme_options = {
+    #    'canonical_url': '',
+    #    'analytics_id': 'UA-XXXXXXX-1',  #  Provided by Google in your dashboard
+    "logo_only": True,
+    "display_version": True,
+    "prev_next_buttons_location": "both",
+    "style_external_links": False,
+    "style_nav_header_background": "#3c4142",
+    # Toc options
+    "collapse_navigation": True,
+    "sticky_navigation": True,
+    "navigation_depth": 2,
+    "includehidden": True,
+    "titles_only": False,
+}
+
+# Add any paths that contain custom static files (such as style sheets) here,
+# 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"]
+html_css_files = ["custom.css"]
+
+# The name of an image file (relative to this directory) to place at the top
+# of the sidebar.
+# html_logo = "_static/images/logo_small_clear.png"
+
+# The name of an image file (within the static path) to use as favicon of the
+# docs.  This file should be a Windows icon file (.ico) being 16x16 or 32x32
+# pixels large.
+# html_favicon = "_static/images/desc_icon.ico"
+
+
+# If not '', a 'Last updated on:' timestamp is inserted at every page bottom,
+# using the given strftime format.
+html_last_updated_fmt = "%b %d, %Y"
+
+# If true, SmartyPants will be used to convert quotes and dashes to
+# typographically correct entities.
+html_use_smartypants = True
+
+# Custom sidebar templates, maps document names to template names.
+# html_sidebars = {}
+
+# Additional templates that should be rendered to pages, maps page names to
+# template names.
+# html_additional_pages = {}
+
+# If false, no module index is generated.
+html_domain_indices = True
+
+# If false, no index is generated.
+html_use_index = True
+
+# If true, the index is split into individual pages for each letter.
+html_split_index = False
+
+# If true, links to the reST sources are added to the pages.
+html_show_sourcelink = True
+
+# If true, "Created using Sphinx" is shown in the HTML footer. Default is True.
+html_show_sphinx = True
+
+# If true, "(C) Copyright ..." is shown in the HTML footer. Default is True.
+html_show_copyright = True
+
+# If true, an OpenSearch description file will be output, and all pages will
+# contain a <link> tag referring to it.  The value of this option must be the
+# base URL from which the finished HTML is served.
+# html_use_opensearch = ''
+
+# This is the file name suffix for HTML files (e.g. ".xhtml").
+# html_file_suffix = None
+
+# Output file base name for HTML help builder.
+htmlhelp_basename = "tsadar"
diff --git a/docs/source/fitmodel.rst b/docs/source/fitmodel.rst
new file mode 100644
index 0000000..2fd78ce
--- /dev/null
+++ b/docs/source/fitmodel.rst
@@ -0,0 +1,10 @@
+FitModel API
+================
+
+This is the module that handles the loss function, gradient calculation, hessian calculation, and parameter generation
+
+.. autoclass:: inverse_thomson_scattering.model.physics.generate_spectra.FitModel
+    :members:
+    :private-members:
+
+
diff --git a/docs/source/fitter.rst b/docs/source/fitter.rst
new file mode 100644
index 0000000..260d27f
--- /dev/null
+++ b/docs/source/fitter.rst
@@ -0,0 +1,9 @@
+Fitter API
+============
+
+This is the highest level module that handles the fitting procedure
+
+.. automodule:: inverse_thomson_scattering.fitter
+    :members:
+    :private-members:
+    :special-members:
\ No newline at end of file
diff --git a/docs/source/fitting.rst b/docs/source/fitting.rst
index 3d7b8ae..6a61d6b 100644
--- a/docs/source/fitting.rst
+++ b/docs/source/fitting.rst
@@ -1,7 +1,8 @@
 Best practices for fitting
 ---------------------------------
 
-It is recommended start fitting data with a coarse resolution in order to identify the rough plasma conditions. These conditions can then be used as the initial conditions for a fine resolution fit.
+It is recommended start fitting data with a coarse resolution in order to identify the rough plasma conditions.
+These conditions can then be used as the initial conditions for a fine resolution fit.
 
 
 
@@ -9,7 +10,8 @@ It is recommended start fitting data with a coarse resolution in order to identi
 Background and lineout selection
 ---------------------------------
 
-There are multiple options for background algorithms and types of fitting. These tend to be the best options for various data types. All of these options are editable in the input deck.
+There are multiple options for background algorithms and types of fitting. These tend to be the best options for various
+data types. All of these options are editable in the input deck.
 
 **Best operation for time resolved data:**
 
diff --git a/docs/source/formfactor.rst b/docs/source/formfactor.rst
new file mode 100644
index 0000000..8789315
--- /dev/null
+++ b/docs/source/formfactor.rst
@@ -0,0 +1,11 @@
+Form Factor API
+================
+
+This is the module that handles the loss function, gradient calculation, hessian calculation, and parameter generation
+
+.. autoclass:: inverse_thomson_scattering.model.physics.form_factor.FormFactor
+    :members:
+    :private-members:
+
+.. autofunction:: inverse_thomson_scattering.model.physics.form_factor.zprimeMaxw
+
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 072e1bb..c0c33f1 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -26,6 +26,7 @@ More details on the numerics will be added soon.
    FAQ
    math
    contributing
+   api_main
    
 *What does TSADAR stand for?*
 
diff --git a/docs/source/math.rst b/docs/source/math.rst
index 6b6a8e6..e1afc3c 100644
--- a/docs/source/math.rst
+++ b/docs/source/math.rst
@@ -1,4 +1,4 @@
 Implementation of Thomson Scattering
----------------------------------
+-------------------------------------
 
 *Coming Soon*
\ No newline at end of file
diff --git a/docs/source/ratintn.rst b/docs/source/ratintn.rst
new file mode 100644
index 0000000..46dbdc5
--- /dev/null
+++ b/docs/source/ratintn.rst
@@ -0,0 +1,10 @@
+Dispersion Relation Integral
+=============================
+
+The dispersion relation integral is the most computationally expensive component. Because there is a pole in the
+integral, we perform a rational integral that is described by the following routines
+
+
+.. autofunction::   inverse_thomson_scattering.model.physics.ratintn.ratintn
+.. autofunction::   inverse_thomson_scattering.model.physics.ratintn.ratcen
+
diff --git a/docs/source/runner.rst b/docs/source/runner.rst
new file mode 100644
index 0000000..255b943
--- /dev/null
+++ b/docs/source/runner.rst
@@ -0,0 +1,9 @@
+Runner API
+============
+
+This is the highest level module that manages the MLFlow run and runs the calculation
+
+.. automodule:: inverse_thomson_scattering.runner
+    :members:
+    :private-members:
+    :special-members:
\ No newline at end of file
diff --git a/docs/source/spectrum.rst b/docs/source/spectrum.rst
new file mode 100644
index 0000000..9a533cc
--- /dev/null
+++ b/docs/source/spectrum.rst
@@ -0,0 +1,8 @@
+Spectrum API
+================
+
+TODO
+
+.. autoclass:: inverse_thomson_scattering.model.spectrum.SpectrumCalculator
+    :members:
+    :private-members:
diff --git a/env.yml b/env.yml
index e540a84..39f31dd 100644
--- a/env.yml
+++ b/env.yml
@@ -5,8 +5,8 @@ dependencies:
   - python=3.10
   - pip
   - pip:
-      - jaxlib==0.4.14
-      - jax==0.4.14
+      - jaxlib
+      - jax
       - jaxopt
       - numpy
       - scipy
@@ -20,4 +20,6 @@ dependencies:
       - typing-extensions
       - optax
       - tqdm
-      - xarray
\ No newline at end of file
+      - xarray
+      - mlflow_export_import
+      - pandas
\ No newline at end of file
diff --git a/env_gpu.yml b/env_gpu.yml
index 2f41cd0..da6b778 100644
--- a/env_gpu.yml
+++ b/env_gpu.yml
@@ -25,4 +25,5 @@ dependencies:
       - optax
       - tqdm
       - xarray
-      - mlflow_export_import
\ No newline at end of file
+      - mlflow_export_import
+      - pandas
\ No newline at end of file
diff --git a/inverse_thomson_scattering/fitter.py b/inverse_thomson_scattering/fitter.py
index b060c18..270cd46 100644
--- a/inverse_thomson_scattering/fitter.py
+++ b/inverse_thomson_scattering/fitter.py
@@ -1,5 +1,5 @@
 import copy
-from typing import Dict, Tuple
+from typing import Dict, Tuple, List
 import time
 import numpy as np
 import pandas as pd
@@ -10,7 +10,7 @@
 from jax.flatten_util import ravel_pytree
 
 from inverse_thomson_scattering.misc.num_dist_func import get_num_dist_func
-from inverse_thomson_scattering.model.loss_function import TSFitter
+from inverse_thomson_scattering.model.TSFitter import TSFitter
 from inverse_thomson_scattering.process import prepare, postprocess
 
 
@@ -102,7 +102,7 @@ def _validate_inputs_(config: Dict) -> Dict:
     return config
 
 
-def scipy_angular_loop(config: Dict, all_data: Dict, sa):
+def scipy_angular_loop(config: Dict, all_data: Dict, sa) -> Tuple[Dict, float, TSFitter]:
     """
     Performs angular thomson scattering i.e. ARTEMIS fitting exercise using the SciPy optimizer routines
 
@@ -134,7 +134,7 @@ def scipy_angular_loop(config: Dict, all_data: Dict, sa):
 
     ts_fitter = TSFitter(config, sa, batch)
     all_weights = {k: [] for k in ts_fitter.pytree_weights["active"].keys()}
-    
+
     if config["optimizer"]["num_mins"] > 1:
         print(Warning("multiple num mins doesnt work. only running once"))
     for i in range(1):
@@ -248,9 +248,11 @@ def angular_adam(config, all_data, sa, batch_indices, num_batches):
             best_weights = weights
 
         mlflow.log_metrics({"epoch loss": float(epoch_loss)}, step=i_epoch)
-        
 
-def _1d_adam_loop_(config, ts_fitter, previous_weights, batch, tbatch):
+
+def _1d_adam_loop_(
+    config: Dict, ts_fitter: TSFitter, previous_weights: np.ndarray, batch: Dict, tbatch
+) -> Tuple[float, Dict]:
     jaxopt_kwargs = dict(
         fun=ts_fitter.vg_loss, maxiter=config["optimizer"]["num_epochs"], value_and_grad=True, has_aux=True
     )
@@ -283,8 +285,9 @@ def _1d_adam_loop_(config, ts_fitter, previous_weights, batch, tbatch):
             best_weights = init_weights
 
     return best_loss, best_weights
-  
-def _1d_scipy_loop_(config, ts_fitter: TSFitter, previous_weights, batch, tbatch):
+
+
+def _1d_scipy_loop_(config: Dict, ts_fitter: TSFitter, previous_weights: np.ndarray, batch: Dict) -> Tuple[float, Dict]:
     if previous_weights is None:  # if prev, then use that, if not then use flattened weights
         init_weights = np.copy(ts_fitter.flattened_weights)
     else:
@@ -304,14 +307,33 @@ def _1d_scipy_loop_(config, ts_fitter: TSFitter, previous_weights, batch, tbatch
         bounds=ts_fitter.bounds,
         options={"disp": True, "maxiter": 1000},
     )
-    
+
     best_loss = res["fun"]
     best_weights = ts_fitter.unravel_pytree(res["x"])
 
     return best_loss, best_weights
 
 
-def one_d_loop(config, all_data, sa, batch_indices, num_batches):
+def one_d_loop(
+    config: Dict, all_data: Dict, sa: Tuple, batch_indices: np.ndarray, num_batches: int
+) -> Tuple[List, float, TSFitter]:
+    """
+    This is the higher level wrapper that prepares the data and the fitting code for the 1D fits
+
+    This function branches out into the various optimization routines for fitting.
+
+    For now, this is either running the ADAM loop or the SciPy optimizer loop
+
+    Args:
+        config:
+        all_data:
+        sa:
+        batch_indices:
+        num_batches:
+
+    Returns:
+
+    """
     sample = {k: v[: config["optimizer"]["batch_size"]] for k, v in all_data.items()}
     sample = {
         "noise_e": config["other"]["PhysParams"]["noiseE"][: config["optimizer"]["batch_size"]],
@@ -342,7 +364,7 @@ def one_d_loop(config, all_data, sa, batch_indices, num_batches):
             else:
                 # not sure why this is needed but something needs to be reset, either the weights or the bounds
                 ts_fitter = TSFitter(config, sa, batch)
-                best_loss, best_weights = _1d_scipy_loop_(config, ts_fitter, previous_weights, batch, tbatch)
+                best_loss, best_weights = _1d_scipy_loop_(config, ts_fitter, previous_weights, batch)
 
             all_weights.append(best_weights)
             mlflow.log_metrics({"batch loss": float(best_loss)}, step=i_batch)
diff --git a/inverse_thomson_scattering/model/loss_function.py b/inverse_thomson_scattering/model/TSFitter.py
similarity index 80%
rename from inverse_thomson_scattering/model/loss_function.py
rename to inverse_thomson_scattering/model/TSFitter.py
index a0e2eb9..8062cfe 100644
--- a/inverse_thomson_scattering/model/loss_function.py
+++ b/inverse_thomson_scattering/model/TSFitter.py
@@ -13,7 +13,23 @@
 
 
 class TSFitter:
+    """
+    This class is responsible for handling the forward pass and using that to create a loss function
+
+    Args:
+            cfg: Configuration dictionary
+            sas: TODO
+            dummy_batch: Dictionary of dummy data
+
+    """
     def __init__(self, cfg: Dict, sas, dummy_batch):
+        """
+
+        Args:
+            cfg: Configuration dictionary
+            sas: TODO
+            dummy_batch: Dictionary of dummy data
+        """
         self.cfg = cfg
 
         if cfg["optimizer"]["y_norm"]:
@@ -64,9 +80,24 @@ def __init__(self, cfg: Dict, sas, dummy_batch):
             else:
                 raise NotImplementedError
         else:
-            Warning("\n !!! Distribution function not fitted !!! Make sure this is what you thought you were running \n")
+            Warning(
+                "\n !!! Distribution function not fitted !!! Make sure this is what you thought you were running \n"
+            )
+
+    def smooth(self, distribution: jnp.ndarray) -> jnp.ndarray:
+        """
+        This method is used to smooth the distribution function. It sits right in between the optimization algorithm
+        that provides the weights/values of the distribution function and the fitting code that uses it.
 
-    def smooth(self, distribution):
+        Because the optimizer is not constrained to provide a smooth distribution function, this operation smoothens
+        the output. This is a differentiable operation and we train/fit our weights through this.
+
+        Args:
+            distribution:
+
+        Returns:
+
+        """
         s = jnp.r_[
             distribution[self.smooth_window_len - 1 : 0 : -1],
             distribution,
@@ -76,7 +107,20 @@ def smooth(self, distribution):
             self.smooth_window_len - 1 : -(self.smooth_window_len - 1)
         ]
 
-    def weights_to_params(self, these_params, return_static_params=True):
+    def weights_to_params(self, these_params: Dict, return_static_params: bool = True) -> Dict:
+        """
+        This function creates the physical parameters used in the TS algorithm from the weights.
+
+        This could be a 1:1 mapping, or it could be a linear transformation e.g. "normalized" parameters, or it could
+        be something else altogether e.g. a neural network
+
+        Args:
+            these_params:
+            return_static_params:
+
+        Returns:
+
+        """
         for param_name, param_config in self.cfg["parameters"].items():
             if param_config["active"]:
                 # if self.cfg["optimizer"]["method"] == "adam":
@@ -95,7 +139,19 @@ def weights_to_params(self, these_params, return_static_params=True):
 
         return these_params
 
-    def _array_loss_fn_(self, weights, batch):
+    def _array_loss_fn_(self, weights: Dict, batch: Dict):
+        """
+        This is similar to the _loss_ function but it calculates the loss per slice without doing a batch-wise mean
+        calculation at the end. There might be a better way to write this
+
+
+        Args:
+            weights:
+            batch:
+
+        Returns:
+
+        """
         # Used for postprocessing
         params = self.weights_to_params(weights)
         ThryE, ThryI, lamAxisE, lamAxisI = self.spec_calc(params, batch)
@@ -212,7 +268,17 @@ def calc_ei_error(self, batch, ThryI, lamAxisI, ThryE, lamAxisE, denom, reduce_f
 
         return i_error, e_error
 
-    def moment_loss(self, params):
+    def _moment_loss_(self, params):
+        """
+        This function calculates the loss associated with regularizing the moments of the distribution function i.e.
+        the density should be 1, the temperature should be 1, and momentum should be 0.
+
+        Args:
+            params:
+
+        Returns:
+
+        """
         dv = self.cfg["velocity"][1] - self.cfg["velocity"][0]
         if self.cfg["parameters"]["fe"]["symmetric"]:
             density_loss = jnp.mean(jnp.square(1.0 - 2.0 * jnp.sum(jnp.exp(params["fe"]) * dv, axis=1)))
@@ -264,16 +330,25 @@ def __loss__(self, weights, batch):
             denom=[jnp.square(self.i_norm), jnp.square(self.e_norm)],
             reduce_func=jnp.mean,
         )
-        density_loss, temperature_loss, momentum_loss = self.moment_loss(params)
+        density_loss, temperature_loss, momentum_loss = self._moment_loss_(params)
         # other_losses = calc_other_losses(params)
-        normed_batch = self.get_normed_batch(batch)
+        normed_batch = self._get_normed_batch_(batch)
         normed_e_data = normed_batch["e_data"]
         ion_error = self.cfg["data"]["ion_loss_scale"] * i_error
 
         total_loss = ion_error + e_error + density_loss + temperature_loss + momentum_loss
         return total_loss, [ThryE, normed_e_data, params]
 
-    def get_normed_batch(self, batch):
+    def _get_normed_batch_(self, batch):
+        """
+        Normalizes the batch
+
+        Args:
+            batch:
+
+        Returns:
+
+        """
         normed_batch = copy.deepcopy(batch)
         normed_batch["i_data"] = normed_batch["i_data"] / self.i_input_norm
         normed_batch["e_data"] = normed_batch["e_data"] / self.e_input_norm
@@ -287,7 +362,22 @@ def loss(self, weights, batch):
         else:
             return self._loss_(weights, batch)
 
-    def vg_loss(self, weights, batch):
+    def vg_loss(self, weights: Dict, batch: Dict):
+        """
+        This is the primary workhorse high level function. This function returns the value of the loss function which
+        is used to assess goodness-of-fit and the gradient of that value with respect to the weights, which is used to
+        update the weights
+
+        This function is used by both optimization methods. It performs the necessary pre-/post- processing that is
+        needed to work with the optimization software.
+
+        Args:
+            weights:
+            batch:
+
+        Returns:
+
+        """
         if self.cfg["optimizer"]["method"] == "l-bfgs-b":
             pytree_weights = self.unravel_pytree(weights)
             (value, aux), grad = self._vg_func_(pytree_weights, batch)
diff --git a/inverse_thomson_scattering/model/physics/form_factor.py b/inverse_thomson_scattering/model/physics/form_factor.py
index ef176d0..6d571ec 100644
--- a/inverse_thomson_scattering/model/physics/form_factor.py
+++ b/inverse_thomson_scattering/model/physics/form_factor.py
@@ -54,7 +54,12 @@ def __init__(self, lamrang, npts):
 
     def __call__(self, params, cur_ne, cur_Te, sa, f_and_v, lam):
         """
-        NONMAXWTHOMSON calculates the Thomson spectral density function S(k,omg) and is capable of handeling multiple plasma conditions and scattering angles. The spectral density function is calculated with and without the ion contribution which can be set to an independent grid from the electron contribution. Distribution functions can be one or two dimensional and the appropriate susceptibility is calculated with the rational integration.
+        Calculates the Thomson spectral density function S(k,omg) and is capable of handeling multiple plasma conditions
+        and scattering angles. The spectral density function is calculated with and without the ion contribution
+        which can be set to an independent grid from the electron contribution. Distribution functions can be one or
+        two dimensional and the appropriate susceptibility is calculated with the rational integration.
+
+        In angular, `fe` is a Tuple, Distribution function (DF), normalized velocity (x), and angles from k_L to f1 in radians
 
 
 
@@ -69,8 +74,8 @@ def __call__(self, params, cur_ne, cur_Te, sa, f_and_v, lam):
         :param lamrang: wavelength range in nm [1 by 2]
         :param lam: probe wavelength in nm
         :param sa: scattering angle in degrees [1 by n]
-        :param fe: Distribution function (DF) and normalized velocity (x) for 1D distributions and
-        Distribution function (DF), normalized velocity (x), and angles from k_L to f1 in radians
+        :param fe: Distribution function (DF) and normalized velocity (x) for 1D distributions
+
         :return:
         """
 
diff --git a/inverse_thomson_scattering/model/physics/generate_spectra.py b/inverse_thomson_scattering/model/physics/generate_spectra.py
index 034161b..767725e 100644
--- a/inverse_thomson_scattering/model/physics/generate_spectra.py
+++ b/inverse_thomson_scattering/model/physics/generate_spectra.py
@@ -1,3 +1,5 @@
+from typing import Dict
+
 from inverse_thomson_scattering.model.physics.form_factor import FormFactor
 from inverse_thomson_scattering.misc.num_dist_func import get_num_dist_func
 
@@ -5,14 +7,32 @@
 
 
 class FitModel:
-    def __init__(self, config, sa):
+    """
+    TODO
+
+    Args:
+        config:
+        sa:
+    """
+
+    def __init__(self, config: Dict, sa):
         self.config = config
         self.sa = sa
         self.electron_form_factor = FormFactor(config["other"]["lamrangE"], npts=config["other"]["npts"])
         self.ion_form_factor = FormFactor(config["other"]["lamrangI"], npts=config["other"]["npts"])
         self.num_dist_func = get_num_dist_func(config["parameters"]["fe"]["type"], config["velocity"])
 
-    def __call__(self, all_params):
+    def __call__(self, all_params: Dict):
+        """
+        TODO
+
+
+        Args:
+            all_params:
+
+        Returns:
+
+        """
         for key in self.config["parameters"].keys():
             if key != "fe":
                 all_params[key] = jnp.squeeze(all_params[key])
diff --git a/inverse_thomson_scattering/model/spectrum.py b/inverse_thomson_scattering/model/spectrum.py
index ef569cc..3df5d5a 100644
--- a/inverse_thomson_scattering/model/spectrum.py
+++ b/inverse_thomson_scattering/model/spectrum.py
@@ -4,9 +4,24 @@
 from inverse_thomson_scattering.model.physics.generate_spectra import FitModel
 from inverse_thomson_scattering.process import irf
 
+"""
+This is the class that generates the spectrum (?)
+
+TODO 
+
+"""
 
 class SpectrumCalculator:
+    """
+    TODO
+
+    Args:
+        cfg:
+        sas:
+        dummy_batch:
+    """
     def __init__(self, cfg, sas, dummy_batch):
+
         super().__init__()
         self.cfg = cfg
         self.sas = sas
@@ -26,6 +41,20 @@ def __init__(self, cfg, sas, dummy_batch):
             self.vmap_postprocess_thry = vmap(self.postprocess_thry)
 
     def postprocess_thry(self, modlE, modlI, lamAxisE, lamAxisI, amps, TSins):
+        """
+        TODO
+
+        Args:
+            modlE:
+            modlI:
+            lamAxisE:
+            lamAxisI:
+            amps:
+            TSins:
+
+        Returns:
+
+        """
         if self.cfg["other"]["extraoptions"]["load_ion_spec"]:
             lamAxisI, ThryI = irf.add_ion_IRF(self.cfg, lamAxisI, modlI, amps["i_amps"], TSins)
         else:
@@ -45,6 +74,18 @@ def postprocess_thry(self, modlE, modlI, lamAxisE, lamAxisI, amps, TSins):
         return ThryE, ThryI, lamAxisE, lamAxisI
 
     def reduce_ATS_to_resunit(self, ThryE, lamAxisE, TSins, batch):
+        """
+        TODO
+
+        Args:
+            ThryE:
+            lamAxisE:
+            TSins:
+            batch:
+
+        Returns:
+
+        """
         lam_step = round(ThryE.shape[1] / batch["e_data"].shape[1])
         ang_step = round(ThryE.shape[0] / self.cfg["other"]["CCDsize"][0])
 
@@ -60,6 +101,16 @@ def reduce_ATS_to_resunit(self, ThryE, lamAxisE, TSins, batch):
         return ThryE, lamAxisE
 
     def __call__(self, params, batch):
+        """
+        TODO
+
+        Args:
+            params:
+            batch:
+
+        Returns:
+
+        """
         modlE, modlI, lamAxisE, lamAxisI, live_TSinputs = self.vmap_forward_pass(params)  # , sas["weights"])
         ThryE, ThryI, lamAxisE, lamAxisI = self.vmap_postprocess_thry(
             modlE, modlI, lamAxisE, lamAxisI, {"e_amps": batch["e_amps"], "i_amps": batch["i_amps"]}, live_TSinputs
diff --git a/inverse_thomson_scattering/process/correct_throughput.py b/inverse_thomson_scattering/process/correct_throughput.py
index f3f353b..32c6535 100644
--- a/inverse_thomson_scattering/process/correct_throughput.py
+++ b/inverse_thomson_scattering/process/correct_throughput.py
@@ -3,26 +3,39 @@
 from numpy.matlib import repmat
 import scipy.interpolate as sp
 import scipy.io as sio
-from os.path import join, exists
-import matplotlib.pyplot as plt
+from os.path import join
 
 
 def correctThroughput(data, tstype, axisy, shotNum):
+    """
+    Corrects throughput
+
+    todo: improve documentation and typehints
+
+    Args:
+        data:
+        tstype:
+        axisy:
+        shotNum:
+
+    Returns:
+
+    """
     if tstype == "angular":
-        imp = sio.loadmat(join('files', 'spectral_sensitivity.mat'), variable_names='speccal')
-        speccal = imp['speccal']
+        imp = sio.loadmat(join("files", "spectral_sensitivity.mat"), variable_names="speccal")
+        speccal = imp["speccal"]
         # not sure if this transpose is correct, need to check once plotting is working
         speccal = np.transpose(speccal)
         if shotNum < 95000:
-            vq1 = 1. / speccal
+            vq1 = 1.0 / speccal
         else:
-            specax = np.arange(0, 1023) * .214116 + 449.5272
+            specax = np.arange(0, 1023) * 0.214116 + 449.5272
             speccalshift = sp.interp1d(specax, speccal, "linear", bounds_error=False, fill_value=speccal[0])
-            vq1 = 1. / speccalshift(axisy)
+            vq1 = 1.0 / speccalshift(axisy)
         # print(np.shape(vq1))
 
     elif tstype == "temporal":
-        wb = xlrd.open_workbook(join('files', 'Copy of MeasuredSensitivity_9.21.15.xls'))
+        wb = xlrd.open_workbook(join("files", "Copy of MeasuredSensitivity_9.21.15.xls"))
         sheet = wb.sheet_by_index(0)
         sens = np.zeros([301, 2])
 
@@ -30,7 +43,7 @@ def correctThroughput(data, tstype, axisy, shotNum):
             sens[i - 2, 0] = sheet.cell_value(i, 0)
             sens[i - 2, 1] = sheet.cell_value(i, 1)
 
-        sens[:, 1] = 1. / sens[:, 1]
+        sens[:, 1] = 1.0 / sens[:, 1]
         sens[0:17, 1] = sens[18, 1]  # the sensitivity goes to zero in this location and is not usable
 
         speccalshift = sp.interp1d(sens[:, 0], sens[:, 1], "linear", bounds_error=False, fill_value=sens[0, 1])
diff --git a/inverse_thomson_scattering/process/evaluate_background.py b/inverse_thomson_scattering/process/evaluate_background.py
index ac11f98..79ed1c8 100644
--- a/inverse_thomson_scattering/process/evaluate_background.py
+++ b/inverse_thomson_scattering/process/evaluate_background.py
@@ -1,4 +1,5 @@
-# Contains two functions one which quantifies the background for full images and one whcih quantifies the background for individual lineouts
+from typing import Tuple
+
 import numpy as np
 import scipy.optimize as spopt
 import matplotlib.pyplot as plt
@@ -10,6 +11,17 @@
 
 
 def get_shot_bg(config, axisyE, elecData):
+    """
+    Quantify the background for full images
+
+    Args:
+        config:
+        axisyE:
+        elecData:
+
+    Returns:
+
+    """
     if config["data"]["background"]["type"] == "Shot":
         [BGele, BGion, _, _] = loadData(
             config["data"]["background"]["slice"], config["data"]["shotDay"], config["other"]["extraoptions"]
@@ -62,23 +74,32 @@ def quadbg(x):
 
 def get_lineout_bg(
     config, elecData, ionData, BGele, BGion, LineoutTSE_smooth, BackgroundPixel, LineoutPixelE, LineoutPixelI
-):
+) -> Tuple[np.ndarray, np.ndarray]:
     """
-    This function generates noise or background profiles to based off the data or background data. Electron spectra have 2 options "Fit" and "pixel". These specify how forground data is treated. Noise is then the sum of forground and background noise. Ions only have one background option as the background is usualy very small
-    
-    "Fit" : fits a rational model against the edges of the lineout to produce a background, the model can be changed. This option functions differently for angular data and is handled by the function get_shot_bg. This option makes no attempt to remove a background shot, using both can result in double counting. This option is best for imaging data.
-    
-    "pixel : the other options "ps" and "auto" are aliases for "pixel" where the background pixel is instead idenitifed by a time ("ps") or set to 100 pixels past the lineout ("auto"). This method uses another linout of the data that is smoothed to act as the background. If included a background shot is removed to prevent double counting. This option is best for time resolved data.
+    This function generates noise or background profiles to based off the data or background data.
+    Electron spectra have 2 options "Fit" and "pixel". These specify how foreground data is treated.
+    Noise is then the sum of foreground and background noise. Ions only have one background option as the background is
+    usually very small
+
+    "Fit" : fits a rational model against the edges of the lineout to produce a background, the model can be changed.
+    This option functions differently for angular data and is handled by the function get_shot_bg. This option makes no
+    attempt to remove a background shot, using both can result in double counting. This option is best for imaging data.
+
+    "pixel : the other options "ps" and "auto" are aliases for "pixel" where the background pixel is instead identified
+    by a time ("ps") or set to 100 pixels past the lineout ("auto"). This method uses another lint of the data that is
+     smoothed to act as the background. If included a background shot is removed to prevent double counting. This option
+      is best for time resolved data.
     """
     span = 2 * config["data"]["dpixel"] + 1  # (span must be odd)
 
     if config["data"]["background"]["type"] not in ["Fit", "Shot", "pixel"]:
         raise NotImplementedError("Background type must be: 'Fit', 'Shot', or 'pixel'")
-    
+
     if config["other"]["extraoptions"]["load_ele_spec"]:
         if config["data"]["background"]["type"] == "Fit":
             if config["other"]["extraoptions"]["spectype"] != "angular":
-                # exp2 bg seems to be the best for some imaging data while rat11 is better in other cases but should be checked in more situations
+                # exp2 bg seems to be the best for some imaging data while rat11 is better in other cases but
+                # should be checked in more situations
                 bgfitx = np.hstack([np.arange(100, 200), np.arange(800, 1023)])
 
                 def exp2(x, a, b, c, d):
@@ -106,9 +127,9 @@ def rat11(x, a, b, c):
                         plt.plot(rat11(np.arange(1024), *rat1bg))
                         plt.plot(LineoutTSE_smooth[i])
                         plt.show()
-                    
+
                     LineoutBGE.append(rat11(np.arange(1024), *rat1bg))
-        #if not fit
+        # if not fit
         else:
             # quantify a background lineout
             LineoutBGE = np.mean(
@@ -130,7 +151,7 @@ def exp2(x, a, b, c, d):
                 )  # this is specificaly targeted at streaked data, removes the fiducials at top and bottom and notch filter
                 bgfitx2 = np.hstack([np.arange(250, 300), np.arange(700, 900)])
                 plt.plot(bgfitx, LineoutBGE[bgfitx])
-                
+
                 [expbg, _] = spopt.curve_fit(exp2, bgfitx, LineoutBGE[bgfitx], p0=[200, 0.001, 200, 0.001])
                 LineoutBGE = config["data"]["bgscaleE"] * exp2(np.arange(1024), *expbg)
 
@@ -152,7 +173,7 @@ def exp2(x, a, b, c, d):
                 #     plt.plot(bgfitx, lin[bgfitx])
                 #     plt.show()
 
-        #add background from background shot is applicable
+        # add background from background shot is applicable
         if np.shape(BGele) == tuple(config["other"]["CCDsize"]):
             LineoutBGE2 = [
                 np.mean(BGele[:, a - config["data"]["dpixel"] : a + config["data"]["dpixel"]], axis=1)
@@ -161,19 +182,18 @@ def exp2(x, a, b, c, d):
             noiseE = LineoutBGE + np.array(LineoutBGE2)
         else:
             noiseE = LineoutBGE * np.ones((len(LineoutPixelE), 1))
-            
+
         # constant addition to the background
         noiseE += config["other"]["flatbg"]
 
     else:
         noiseE = np.zeros(len(config["data"]["lineouts"]["val"]))
-    
-    
+
     if config["other"]["extraoptions"]["load_ion_spec"]:
-        #Due to the low background associated with IAWs the fitted background is only performed for the EPW
+        # Due to the low background associated with IAWs the fitted background is only performed for the EPW
         if config["data"]["background"]["type"] == "Fit":
             BackgroundPixel = config["data"]["background"]["slice"]
-        
+
         # quantify a uniform background
         noiseI = np.mean(
             (ionData - BGion)[
@@ -189,9 +209,7 @@ def exp2(x, a, b, c, d):
         # add the uniform background to the background from the background shot
         if np.shape(BGion) == tuple(config["other"]["CCDsize"]):
             LineoutBGI = [
-                np.mean(
-                    BGion[:, a - config["data"]["dpixel"] : a + config["data"]["dpixel"]], axis=1
-                )
+                np.mean(BGion[:, a - config["data"]["dpixel"] : a + config["data"]["dpixel"]], axis=1)
                 for a in LineoutPixelI
             ]
             noiseI = noiseI + LineoutBGI
diff --git a/inverse_thomson_scattering/process/irf.py b/inverse_thomson_scattering/process/irf.py
index 065906e..d85d1dd 100644
--- a/inverse_thomson_scattering/process/irf.py
+++ b/inverse_thomson_scattering/process/irf.py
@@ -1,7 +1,25 @@
+from typing import Tuple
 from jax import numpy as jnp
 
 
-def add_ATS_IRF(config, sas, lamAxisE, modlE, amps, TSins, lam):
+def add_ATS_IRF(config, sas, lamAxisE, modlE, amps, TSins, lam) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """
+    Angular Thomson Scattering IRF
+
+    todo: improve doc and typehints
+
+    Args:
+        config:
+        sas:
+        lamAxisE:
+        modlE:
+        amps:
+        TSins:
+        lam:
+
+    Returns:
+
+    """
     stddev_lam = config["other"]["PhysParams"]["widIRF"]["spect_FWHM_ele"] / 2.3548
     stddev_ang = config["other"]["PhysParams"]["widIRF"]["ang_FWHM_ele"] / 2.3548
     # Conceptual_origin so the convolution donsn't shift the signal
@@ -31,7 +49,22 @@ def add_ATS_IRF(config, sas, lamAxisE, modlE, amps, TSins, lam):
     return lamAxisE, ThryE
 
 
-def add_ion_IRF(config, lamAxisI, modlI, amps, TSins):
+def add_ion_IRF(config, lamAxisI, modlI, amps, TSins) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """
+    Ion IRF (Instrument Response Function?)
+
+    todo: improve doc and typehints
+
+    Args:
+        config:
+        lamAxisI:
+        modlI:
+        amps:
+        TSins:
+
+    Returns:
+
+    """
     stddevI = config["other"]["PhysParams"]["widIRF"]["spect_stddev_ion"]
     if stddevI:
         originI = (jnp.amax(lamAxisI) + jnp.amin(lamAxisI)) / 2.0
@@ -53,7 +86,23 @@ def add_ion_IRF(config, lamAxisI, modlI, amps, TSins):
     return lamAxisI, ThryI
 
 
-def add_electron_IRF(config, lamAxisE, modlE, amps, TSins, lam):
+def add_electron_IRF(config, lamAxisE, modlE, amps, TSins, lam) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """
+    electron IRF (Instrument Response Function?)
+
+    todo: improve doc and typehints
+
+    Args:
+        config:
+        lamAxisE:
+        modlE:
+        amps:
+        TSins:
+        lam:
+
+    Returns:
+
+    """
     stddevE = config["other"]["PhysParams"]["widIRF"]["spect_stddev_ele"]
     # Conceptual_origin so the convolution doesn't shift the signal
     originE = (jnp.amax(lamAxisE) + jnp.amin(lamAxisE)) / 2.0
diff --git a/inverse_thomson_scattering/process/lineouts.py b/inverse_thomson_scattering/process/lineouts.py
index 0f8a4cc..1c96492 100644
--- a/inverse_thomson_scattering/process/lineouts.py
+++ b/inverse_thomson_scattering/process/lineouts.py
@@ -1,3 +1,5 @@
+from typing import Dict
+
 from collections import defaultdict
 
 import numpy as np
@@ -6,7 +8,7 @@
 
 def get_lineouts(
     elecData, ionData, BGele, BGion, axisxE, axisxI, axisyE, axisyI, shift_zero, IAWtime, xlab, sa, config
-):
+) -> Dict:
     # Convert lineout locations to pixel
     if config["data"]["lineouts"]["type"] == "ps" or config["data"]["lineouts"]["type"] == "um":
         LineoutPixelE = [np.argmin(abs(axisxE - loc - shift_zero)) for loc in config["data"]["lineouts"]["val"]]
@@ -65,15 +67,7 @@ def get_lineouts(
 
     # Find background signal combining information from a background shot and background lineout
     [noiseE, noiseI] = get_lineout_bg(
-        config,
-        elecData,
-        ionData,
-        BGele,
-        BGion,
-        LineoutTSE_smooth,
-        BackgroundPixel,
-        LineoutPixelE,
-        LineoutPixelI,
+        config, elecData, ionData, BGele, BGion, LineoutTSE_smooth, BackgroundPixel, LineoutPixelE, LineoutPixelI
     )
 
     # Find data amplitudes
diff --git a/inverse_thomson_scattering/process/postprocess.py b/inverse_thomson_scattering/process/postprocess.py
index d106ab2..07806a1 100644
--- a/inverse_thomson_scattering/process/postprocess.py
+++ b/inverse_thomson_scattering/process/postprocess.py
@@ -6,7 +6,7 @@
 import scipy.optimize as spopt
 
 from inverse_thomson_scattering.misc import plotters
-from inverse_thomson_scattering.model.loss_function import TSFitter
+from inverse_thomson_scattering.model.TSFitter import TSFitter
 
 
 def recalculate_with_chosen_weights(
diff --git a/inverse_thomson_scattering/runner.py b/inverse_thomson_scattering/runner.py
index 67349f5..09212c3 100644
--- a/inverse_thomson_scattering/runner.py
+++ b/inverse_thomson_scattering/runner.py
@@ -1,4 +1,5 @@
 import time, os
+from typing import Dict, Tuple
 
 import numpy as np
 import matplotlib.pyplot as plt
@@ -9,7 +10,7 @@
 from inverse_thomson_scattering import fitter
 from inverse_thomson_scattering.misc.calibration import get_scattering_angles
 from inverse_thomson_scattering.misc.num_dist_func import get_num_dist_func
-from inverse_thomson_scattering.model.loss_function import TSFitter
+from inverse_thomson_scattering.model.TSFitter import TSFitter
 from inverse_thomson_scattering.fitter import init_param_norm_and_shift
 from inverse_thomson_scattering.misc import utils
 
@@ -20,7 +21,7 @@
     BASE_TEMPDIR = None
 
 
-def load_and_make_folders(cfg_path):
+def load_and_make_folders(cfg_path: str) -> Tuple[str, Dict]:
     """
     This is used to queue runs on NERSC
 
@@ -57,7 +58,7 @@ def load_and_make_folders(cfg_path):
     return mlflow_run.info.run_id, all_configs
 
 
-def run(cfg_path, mode):
+def run(cfg_path: str, mode: str) -> str:
     """
     Wrapper for lower level runner
 
@@ -78,7 +79,7 @@ def run(cfg_path, mode):
     return run_id
 
 
-def _run_(config, mode="fit"):
+def _run_(config: Dict, mode: str = "fit"):
     """
     Either performs a forward pass or an entire fitting routine
 
@@ -102,7 +103,7 @@ def _run_(config, mode="fit"):
     mlflow.set_tag("status", "completed")
 
 
-def run_job(run_id, nested):
+def run_job(run_id: str, nested: bool):
     """
     This is used to run queued runs on NERSC. It picks up the `run_id` and finds that using MLFlow and does the fitting
 
@@ -131,7 +132,17 @@ def run_job(run_id, nested):
         utils.export_run(run_id)
 
 
-def calc_spec(config):
+def calc_spec(config: Dict):
+    """
+    Just performs a forward pass
+
+
+    Args:
+        config:
+
+    Returns:
+
+    """
     # get scattering angles and weights
     config["optimizer"]["batch_size"] = 1
     config["other"]["extraoptions"]["spectype"] = "temporal"
@@ -145,11 +156,11 @@ def calc_spec(config):
     config["velocity"] = np.linspace(-7, 7, config["parameters"]["fe"]["length"])
     if config["parameters"]["fe"]["symmetric"]:
         config["velocity"] = np.linspace(0, 7, config["parameters"]["fe"]["length"])
-        
+
     NumDistFunc = get_num_dist_func(config["parameters"]["fe"]["type"], config["velocity"])
     if not config["parameters"]["fe"]["val"]:
         config["parameters"]["fe"]["val"] = np.log(NumDistFunc(config["parameters"]["m"]["val"]))
-    
+
     config["units"] = init_param_norm_and_shift(config)
 
     sas = get_scattering_angles(config)
diff --git a/requirements-docs.txt b/requirements-docs.txt
new file mode 100644
index 0000000..494465f
--- /dev/null
+++ b/requirements-docs.txt
@@ -0,0 +1,7 @@
+-r ./requirements.txt
+
+# building the docs
+sphinx > 3.0.0
+sphinx_copybutton
+sphinx-rtd-theme >= 1.0, < 2.0
+sphinx-github-style >= 1.0, <= 1.1
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 3b45c28..537d977 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -16,3 +16,4 @@ tqdm
 jaxopt
 xarray
 mlflow_export_import
+pandas
\ No newline at end of file