Repository: grburgess/popsynth (language: python)

Issue:
removed stan gen



Discussion / hints:
# [Codecov](https://codecov.io/gh/grburgess/popsynth/pull/7?src=pr&el=h1) Report
> Merging [#7](https://codecov.io/gh/grburgess/popsynth/pull/7?src=pr&el=desc) into [master](https://codecov.io/gh/grburgess/popsynth/commit/caa1db3958729c8b151a967bdc450d35d68c30f5?src=pr&el=desc) will **increase** coverage by `3.55%`.
> The diff coverage is `100%`.

[![Impacted file tree graph](https://codecov.io/gh/grburgess/popsynth/pull/7/graphs/tree.svg?width=650&token=MX0aCZv9t2&height=150&src=pr)](https://codecov.io/gh/grburgess/popsynth/pull/7?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master       #7      +/-   ##
==========================================
+ Coverage      63%   66.55%   +3.55%     
==========================================
  Files          25       25              
  Lines        1619     1507     -112     
==========================================
- Hits         1020     1003      -17     
+ Misses        599      504      -95
```


| [Impacted Files](https://codecov.io/gh/grburgess/popsynth/pull/7?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [popsynth/distributions/pareto\_distribution.py](https://codecov.io/gh/grburgess/popsynth/pull/7/diff?src=pr&el=tree#diff-cG9wc3ludGgvZGlzdHJpYnV0aW9ucy9wYXJldG9fZGlzdHJpYnV0aW9uLnB5) | `62.5% <ø> (+0.73%)` | :arrow_up: |
| [...opsynth/distributions/cosmological\_distribution.py](https://codecov.io/gh/grburgess/popsynth/pull/7/diff?src=pr&el=tree#diff-cG9wc3ludGgvZGlzdHJpYnV0aW9ucy9jb3Ntb2xvZ2ljYWxfZGlzdHJpYnV0aW9uLnB5) | `55.39% <ø> (+7.49%)` | :arrow_up: |
| [popsynth/population.py](https://codecov.io/gh/grburgess/popsynth/pull/7/diff?src=pr&el=tree#diff-cG9wc3ludGgvcG9wdWxhdGlvbi5weQ==) | `81.19% <100%> (+12.86%)` | :arrow_up: |
| [popsynth/\_\_init\_\_.py](https://codecov.io/gh/grburgess/popsynth/pull/7/diff?src=pr&el=tree#diff-cG9wc3ludGgvX19pbml0X18ucHk=) | `100% <0%> (+18.18%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/grburgess/popsynth/pull/7?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/grburgess/popsynth/pull/7?src=pr&el=footer). Last update [caa1db3...fb31ee2](https://codecov.io/gh/grburgess/popsynth/pull/7?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Fix patch:
diff --git a/docs/conf.py b/docs/conf.py
index 28c9290..37b3157 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -47,13 +47,13 @@ def run_apidoc(app):
         str(DOCS / ".." / "popsynth"),
     ])
 
+
 def skip(app, what, name, obj, would_skip, options):
     if name == "__init__":
         return False
     return would_skip
 
 
-    
 # -- General configuration ---------------------------------------------------
 
 # Add any Sphinx extension module names here, as strings. They can be
diff --git a/docs/md/examples.md b/docs/md/examples.md
index c0aa384..30614c6 100644
--- a/docs/md/examples.md
+++ b/docs/md/examples.md
@@ -202,3 +202,136 @@ ax.set_xlabel("log Flux")
 ```
 
 Does this look like the observed catalogs?
+
+## BL Lac blazars
+A model for the luminosity function and cosmic evolution of BL Lac type blazars is presented in [Ajello et al. 2014](https://arxiv.org/abs/1310.0006), based on observations in gamma-rays with the Fermi-LAT instrument.
+
+We can use the results of this paper to build a BL Lac population that is able to reproduce the results reported in the recent [4FGL Fermi-LAT catalog](https://arxiv.org/abs/1902.10045) reasonably well.
+
+```python
+from scipy import special as sf
+from popsynth import (ZPowerCosmoDistribution, SoftFluxSelection,
+                      GalacticPlaneSelection)
+```
+
+The work mentioned above presents 3 models for the BL Lac luminosity function. Here, we focus on the case of pure density evolution (referred to as PDE in the paper). In this case, the BL Lac population is parametrised as having a broken power law luminosity distribution, with an independent density evolution following a cosmological power law distribution.
+
+We work with a luminosity range of $L_\mathrm{min} = 7\times 10^{43}$ erg $\mathrm{s}^{-1}$ and $L_\mathrm{max} = 10^{52}$ erg $\mathrm{s}^{-1}$ following Ajello et al. 2014. All luminosities are in units of erg $\mathrm{s}^{-1}$. Similarly, the maxium redshift considered in $z=6$.
+
+We start by setting up the broken power law distribution (`BPLDistribution`).
+
+```python
+bpl = BPLDistribution()
+bpl.alpha = -1.5
+bpl.beta = -2.5
+bpl.Lmin = 7e43
+bpl.Lmax = 1e52
+bpl.Lbreak = 1e47
+
+fig, ax = plt.subplots()
+L = np.geomspace(bpl.Lmin, bpl.Lmax)
+ax.plot(L, bpl.phi(L), color=purple)
+ax.set_xscale("log")
+ax.set_yscale("log")
+ax.set_xlabel(r"L [erg $\mathrm{s}^{-1}$]")
+```
+
+We now move to the redshift distribution. Following the paper, we parametrize this as a negative power law in $z$. for the purpose of this example, we assume that Bl Lac blazars emit with a steady state. This means that we need to set the `is_rate` parameter to `False` when defining the `ZPowerCosmoDistribution` cosmological distribution. What we mean when we do this is that our local number density, `Lambda` is not per unit time. We also want to survey the whole sky, so we integrate over $4\pi$ sr in the value that we pass to the `Lambda`. 
+
+```python
+zpow = ZPowerCosmoDistribution()
+zpow.Lambda = 9000 # Gpc^-3 sr 
+zpow.delta = -6
+
+fig, ax = plt.subplots()
+z = np.linspace(0.01, 6)
+ax.plot(z, zpow.dNdV(z), color=purple)
+ax.set_yscale("log")
+ax.set_xlabel("z")
+ax.set_ylabel(r"$\frac{\mathrm{d}N}{\mathrm{d}V}$")
+```
+
+Apart from their redshifts and luminosities, BL Lacs also have other properties. As a simple example, we can consider their spectral index, assuming that the gamma-ray emission is well modelled in the energy range of interest (0.1 to 100 GeV) by a simple power law. 
+
+We assume these true values of these indices are normally distributed with mean, $\mu$, and standard deviation, $\tau$. Additionally, we recognise that these are reconstructed quantities in real surveys, with uncertain values. This is reflected in the error, $\sigma$, on the observed values.
+
+```python
+index = NormalAuxSampler(name="index")
+index.mu = 2.1
+index.tau = 0.25
+index.sigma = 0.1
+```
+
+We know that the Fermi-LAT detector cannot detect all objects in the Universe, and it is necessary to model some kind of selection function. In general, brighter and spectrally harder objects are easier to detect. We take this into acount by selecting on the flux, $F=L/4\pi d_L^2(z)$, where $d_L$ is the luminosity distance in cm.
+
+This selection effect will not really be a hard boundary, although we could approximate it as such. In reality, the probability to detect an object increases as a function of its flux. To capture this effect, we consider a `SoftFluxSelection` as follows.
+
+```python
+flux_selector = SoftFluxSelection()
+flux_selector.boundary = 4e-12 # erg cm^-2 s^-1
+flux_selector.strength = 2
+
+# This is what is happening under the hood
+fig, ax = plt.subplots()
+F = np.geomspace(1e-15, 1e-8)
+ax.plot(F, sf.expit(flux_selector.strength * (np.log10(F) - np.log10(flux_selector.boundary))))
+ax.set_xscale("log")
+ax.set_xlabel("F [erg $\mathrm{cm}^{-2}$ $\mathrm{s}^{-1}$]")
+ax.set_ylabel("Detection prob.")
+```
+
+Finally, sometimes it is harder to detect objects near the bright Galactic plane. We take this into account by excluding $10^\circ$ either side of the plane in Galactic longitude using the `GalacticPlaneSelector`.
+
+```python
+gp = GalacticPlaneSelection()
+gp.b_limit = 10
+```
+
+Now, lets finally bring all this together to make a simulated population. Here, we defined our luminosity and spatial distributions already, so we can use them directly in `PopulationSynth`, but there is also the `BPLZPowerCosmoPopulation` available as a quick interface.
+
+```python
+# Main pop synth
+pop_synth = PopulationSynth(spatial_distribution=zpow, luminosity_distribution=bpl)
+
+# Add our selection effects
+pop_synth.set_flux_selection(flux_selector)
+pop_synth.add_spatial_selector(gp)
+
+# Add our auxiliary param - spectral index
+pop_synth.add_observed_quantity(index)
+```
+
+Lets run it! The last parameter to set is adding some uncertainty to our observed flux values.
+
+```python
+population = pop_synth.draw_survey(flux_sigma=0.1)
+```
+
+We can now have a look at the properties of this simulated population, such as the detected and undetected fluxes and distances.
+
+```python
+population.display_fluxes(true_color=purple, obs_color=yellow, with_arrows=False, s=5);
+```
+
+```python
+fig, ax = plt.subplots()
+ax.hist(population.distances, color=purple, histtype="step", lw=3, label="All")
+ax.hist(population.distances[population.selection], histtype="step", lw=3, 
+        color=yellow, label="Detected")
+ax.set_xlabel("$z$")
+ax.legend()
+```
+
+We can also check out the spectral index distribution.
+
+```python
+fig, ax = plt.subplots()
+ax.hist(population.index, color=purple, histtype="step", lw=3, label="All")
+ax.hist(population.index[population.selection], histtype="step", lw=3, 
+        color=yellow, label="Detected")
+ax.set_xlabel("Spectral index")
+ax.legend()
+```
+
+We can now imagine that by changing the input parameters, we can fit our model to the observations in order to have an optimal representation of the true BL Lac blazar population with this parametrisation.
+
diff --git a/popsynth/aux_samplers/__init__.py b/popsynth/aux_samplers/__init__.py
index 18d0714..263e723 100644
--- a/popsynth/aux_samplers/__init__.py
+++ b/popsynth/aux_samplers/__init__.py
@@ -1,11 +1,23 @@
+from .delta_aux_sampler import DeltaAuxSampler
 from .viewing_angle_sampler import ViewingAngleSampler
 from .lognormal_aux_sampler import LogNormalAuxSampler, Log10NormalAuxSampler
 from .normal_aux_sampler import NormalAuxSampler
 from .trunc_normal_aux_sampler import TruncatedNormalAuxSampler
-from .plaw_aux_sampler import ParetoAuxSampler, PowerLawAuxSampler, BrokenPowerLawAuxSampler
 
+from .plaw_aux_sampler import (
+    ParetoAuxSampler,
+    PowerLawAuxSampler,
+    BrokenPowerLawAuxSampler,
+)
 
 __all__ = [
-    "ViewingAngleSampler", "LogNormalAuxSampler", "Log10NormalAuxSampler",
-    "NormalAuxSampler", "TruncatedNormalAuxSampler", "ParetoAuxSampler", "PowerLawAuxSampler", "BrokenPowerLawAuxSampler"
+    "DeltaAuxSampler",
+    "ViewingAngleSampler",
+    "LogNormalAuxSampler",
+    "Log10NormalAuxSampler",
+    "NormalAuxSampler",
+    "TruncatedNormalAuxSampler",
+    "ParetoAuxSampler",
+    "PowerLawAuxSampler",
+    "BrokenPowerLawAuxSampler",
 ]
diff --git a/popsynth/aux_samplers/delta_aux_sampler.py b/popsynth/aux_samplers/delta_aux_sampler.py
index 44fe10f..5fe4657 100644
--- a/popsynth/aux_samplers/delta_aux_sampler.py
+++ b/popsynth/aux_samplers/delta_aux_sampler.py
@@ -11,16 +11,23 @@ class DeltaAuxSampler(AuxiliarySampler):
     sigma = AuxiliaryParameter(default=1, vmin=0)
 
     def __init__(self, name: str, observed: bool = True):
-        """FIXME! briefly describe function
-
-        :param name:
-        :param mu:
-        :param tau:
-        :param sigma:
-        :param observed:
-        :returns:
-        :rtype:
-
+        """
+        A delta-function sampler for which the true value
+        is fixed at ``xp``. Assumes property is observed by default,
+        in which case the observed value is sampled from
+        the true value with some normally-distributed error, ``sigma``.
+
+        :param name: Name of the property
+        :type name: str
+        :param observed: `True` if the property is observed,
+        `False` if it is latent. Defaults to `True`
+        :type observed: bool
+        :param xp: Value at which delta function is located
+        :type xp: :class:`AuxiliaryParameter`
+        :param sigma: Standard deviation of normal distribution
+        from which observed values are sampled, if ``observed``
+        is `True`
+        :type sigma: :class:`AuxiliaryParameter`
         """
 
         super(DeltaAuxSampler, self).__init__(name=name, observed=observed)
diff --git a/popsynth/aux_samplers/lognormal_aux_sampler.py b/popsynth/aux_samplers/lognormal_aux_sampler.py
index 1258309..5b39a6a 100644
--- a/popsynth/aux_samplers/lognormal_aux_sampler.py
+++ b/popsynth/aux_samplers/lognormal_aux_sampler.py
@@ -16,24 +16,29 @@ class LogNormalAuxSampler(AuxiliarySampler):
 
     def __init__(self, name: str, observed: bool = True):
         """
-        A Log normal sampler
-
-        where x ~ e^N(mu, sigma)
-
-
-        :param name:
-        :param observed:
-        :returns:
-        :rtype:
-
+        A Log normal sampler,
+        where property ~ e^N(``mu``, ``sigma``).
+
+        :param name: Name of the property
+        :type name: str
+        :param observed: `True` if the property is observed,
+        `False` if it is latent. Defaults to `True`
+        :type observed: bool
+        :param mu: Mean of the lognormal
+        :type mu: :class:`AuxiliaryParameter`
+        :param tau: Standard deviation of the lognormal
+        :type tau: :class:`AuxiliaryParameter`
+        :param sigma: Standard deviation of normal distribution
+        from which observed values are sampled, if ``observed``
+        is `True`
+        :type sigma: :class:`AuxiliaryParameter`
         """
         super(LogNormalAuxSampler, self).__init__(name=name, observed=observed)
 
     def true_sampler(self, size: int):
 
         self._true_values = np.exp(
-            stats.norm.rvs(loc=self.mu, scale=self.tau,
-                           size=size))  # type: ArrayLike
+            stats.norm.rvs(loc=self.mu, scale=self.tau, size=size))
 
     def observation_sampler(self, size: int):
 
@@ -41,11 +46,11 @@ def observation_sampler(self, size: int):
 
             self._obs_values = stats.norm.rvs(loc=self._true_values,
                                               scale=self.sigma,
-                                              size=size)  # type: ArrayLike
+                                              size=size)
 
         else:
 
-            self._obs_values = self._true_values  # type: ArrayLike
+            self._obs_values = self._true_values
 
 
 class Log10NormalAuxSampler(AuxiliarySampler):
@@ -57,28 +62,30 @@ class Log10NormalAuxSampler(AuxiliarySampler):
 
     def __init__(self, name: str, observed: bool = True):
         """
-        A Log10 normal sampler.
-
-        where x ~ 10^N(mu, sigma)
-
-
-
-        :param name:
-        :param observed:
-        :returns:
-        :rtype:
-
+        A Log10 normal sampler,
+        where property ~ 10^N(``mu``, ``sigma``).
+
+        :param name: Name of the property
+        :type name: str
+        :param observed: `True` if the property is observed,
+        `False` if it is latent. Defaults to `True`
+        :type observed: bool
+        :param mu: Mean of the log10normal
+        :type mu: :class:`AuxiliaryParameter`
+        :param tau: Standard deviation of the log10normal
+        :type tau: :class:`AuxiliaryParameter`
+        :param sigma: Standard deviation of normal distribution
+        from which observed values are sampled, if ``observed``
+        is `True`
+        :type sigma: :class:`AuxiliaryParameter`
         """
         super(Log10NormalAuxSampler, self).__init__(name=name,
                                                     observed=observed)
 
     def true_sampler(self, size: int):
 
-        self._true_values = np.power(10,
-                                     stats.norm.rvs(
-                                         loc=self.mu,
-                                         scale=self.tau,
-                                         size=size))  # type: ArrayLike
+        self._true_values = np.power(
+            10, stats.norm.rvs(loc=self.mu, scale=self.tau, size=size))
 
     def observation_sampler(self, size: int):
 
@@ -86,8 +93,8 @@ def observation_sampler(self, size: int):
 
             self._obs_values = stats.norm.rvs(loc=self._true_values,
                                               scale=self.sigma,
-                                              size=size)  # type: ArrayLike
+                                              size=size)
 
         else:
 
-            self._obs_values = self._true_values  # type: ArrayLike
+            self._obs_values = self._true_values
diff --git a/popsynth/aux_samplers/normal_aux_sampler.py b/popsynth/aux_samplers/normal_aux_sampler.py
index cca2076..8349985 100644
--- a/popsynth/aux_samplers/normal_aux_sampler.py
+++ b/popsynth/aux_samplers/normal_aux_sampler.py
@@ -12,16 +12,24 @@ class NormalAuxSampler(AuxiliarySampler):
 
     def __init__(self, name: str, observed: bool = True):
         """
-        A normal distribution sampler
+        A normal distribution sampler,
+        where property ~ N(``mu``, ``sigma``).
 
-        :param name: 
+        :param name: Name of the property
         :type name: str
-        :param observed: 
+        :param observed: `True` if the property is observed,
+        `False` if it is latent. Defaults to `True`
         :type observed: bool
-        :returns: 
-
+        :param mu: Mean of the normal
+        :type mu: :class:`AuxiliaryParameter`
+        :param tau: Standard deviation of the normal
+        :type tau: :class:`AuxiliaryParameter`
+        :param sigma: Standard deviation of normal distribution
+        from which observed values are sampled, if ``observed``
+        is `True`
+        :type sigma: :class:`AuxiliaryParameter`
         """
-    
+
         super(NormalAuxSampler, self).__init__(name=name, observed=observed)
 
     def true_sampler(self, size: int):
diff --git a/popsynth/aux_samplers/plaw_aux_sampler.py b/popsynth/aux_samplers/plaw_aux_sampler.py
index 178c923..6c1ca78 100644
--- a/popsynth/aux_samplers/plaw_aux_sampler.py
+++ b/popsynth/aux_samplers/plaw_aux_sampler.py
@@ -5,6 +5,7 @@
 from popsynth.auxiliary_sampler import AuxiliarySampler, AuxiliaryParameter
 from popsynth.distributions.bpl_distribution import bpl
 
+
 class ParetoAuxSampler(AuxiliarySampler):
     _auxiliary_sampler_name = "ParetoAuxSampler"
 
@@ -14,21 +15,30 @@ class ParetoAuxSampler(AuxiliarySampler):
 
     def __init__(self, name: str, observed: bool = True):
         """
-        A pareto distribution sampler
+        A pareto distribution sampler,
+        where property ~ 1 / x^(``alpha`` + 1).
 
-        :param name: 
+        :param name: Name of the property
         :type name: str
-        :param observed: 
+        :param observed: `True` if the property is observed,
+        `False` if it is latent. Defaults to `True`
         :type observed: bool
-        :returns: 
-
+        :param xmin: Minimum value of the pareto
+        :type xmin: :class:`AuxiliaryParameter`
+        :param alpha: Index of the pareto
+        :type alpha: :class:`AuxiliaryParameter`
+        :param sigma: Standard deviation of normal distribution
+        from which observed values are sampled, if ``observed``
+        is `True`
+        :type sigma: :class:`AuxiliaryParameter`
         """
-    
+
         super(ParetoAuxSampler, self).__init__(name=name, observed=observed)
 
     def true_sampler(self, size: int):
 
-         self._true_values = (np.random.pareto(self.alpha, size) + 1) * self.xmin
+        self._true_values = (np.random.pareto(self.alpha, size) +
+                             1) * self.xmin
 
     def observation_sampler(self, size: int):
 
@@ -42,6 +52,7 @@ def observation_sampler(self, size: int):
 
             self._obs_values = self._true_values
 
+
 class PowerLawAuxSampler(AuxiliarySampler):
     _auxiliary_sampler_name = "PowerLawAuxSampler"
 
@@ -52,21 +63,30 @@ class PowerLawAuxSampler(AuxiliarySampler):
 
     def __init__(self, name: str, observed: bool = True):
         """
-        A power law distribution sampler
+        A bounded power law distribution sampler,
+        where property ~ x^``alpha``.
 
-        :param name: 
+        :param name: Name of the property
         :type name: str
-        :param observed: 
+        :param observed: `True` if the property is observed,
+        `False` if it is latent. Defaults to `True`
         :type observed: bool
-        :returns: 
-
+        :param xmin: Minimum value of the power law
+        :type xmin: :class:`AuxiliaryParameter`
+        :param xmax: Maximum value of the power law
+        :type xmax: :class:``AuxiliaryParameter
+        :param sigma: Standard deviation of normal distribution
+        from which observed values are sampled, if ``observed``
+        is `True`
+        :type sigma: :class:`AuxiliaryParameter`
         """
-    
+
         super(PowerLawAuxSampler, self).__init__(name=name, observed=observed)
 
     def true_sampler(self, size: int):
 
-        self._true_values = _sample_power_law(self.xmin, self.xmax, self.alpha, size)
+        self._true_values = _sample_power_law(self.xmin, self.xmax, self.alpha,
+                                              size)
 
     def observation_sampler(self, size: int):
 
@@ -90,26 +110,36 @@ class BrokenPowerLawAuxSampler(AuxiliarySampler):
     beta = AuxiliaryParameter()
     xmax = AuxiliaryParameter(vmin=0)
 
-
     def __init__(self, name: str, observed: bool = True):
         """
-        A power law distribution sampler
+        A broken power law distribution sampler,
+        where property ~ x^``alpha`` for x < ``xbreak``,
+        and property ~ x^``beta`` for x > ``xbreak``.
 
-        :param name: 
+        :param name: Name of the property
         :type name: str
-        :param observed: 
+        :param observed: `True` if the property is observed,
+        `False` if it is latent. Defaults to `True`
         :type observed: bool
-        :returns: 
-
+        :param xmin: Minimum value of the broken power law
+        :type xmin: :class:`AuxiliaryParameter`
+        :param xmax: Maximum value of the broken power law
+        :type xmax: :class:``AuxiliaryParameter
+        :param sigma: Standard deviation of normal distribution
+        from which observed values are sampled, if ``observed``
+        is `True`
+        :type sigma: :class:`AuxiliaryParameter`
         """
-    
-        super(BrokenPowerLawAuxSampler, self).__init__(name=name, observed=observed)
+
+        super(BrokenPowerLawAuxSampler, self).__init__(name=name,
+                                                       observed=observed)
 
     def true_sampler(self, size: int):
 
         u = np.atleast_1d(np.random.uniform(size=size))
-        
-        self._true_values = bpl(u, self.xmin, self.xbreak,self.xmax, self.alpha,self.beta)
+
+        self._true_values = bpl(u, self.xmin, self.xbreak, self.xmax,
+                                self.alpha, self.beta)
 
     def observation_sampler(self, size: int):
 
@@ -124,10 +154,6 @@ def observation_sampler(self, size: int):
             self._obs_values = self._true_values
 
 
-
-    
-            
-
 @nb.njit(fastmath=True)
 def _sample_power_law(xmin, xmax, alpha, size):
 
@@ -137,16 +163,11 @@ def _sample_power_law(xmin, xmax, alpha, size):
 
         u = np.random.uniform(0, 1)
         x = np.power(
-            (np.power(xmax, alpha + 1) - np.power(xmin, alpha + 1)) * u
-            + np.power(xmin, alpha + 1),
+            (np.power(xmax, alpha + 1) - np.power(xmin, alpha + 1)) * u +
+            np.power(xmin, alpha + 1),
             1.0 / (alpha + 1.0),
-            )
+        )
 
         out[i] = x
 
-    
     return out
-
-
-
-    
diff --git a/popsynth/aux_samplers/sky_sampler.py b/popsynth/aux_samplers/sky_sampler.py
index 4f9a68c..116229a 100644
--- a/popsynth/aux_samplers/sky_sampler.py
+++ b/popsynth/aux_samplers/sky_sampler.py
@@ -1,79 +1,85 @@
-# import numpy as np
+import numpy as np
 
-# from popsynth.auxiliary_sampler import NonObservedAuxSampler
-# from popsynth.utils.spherical_geometry import sample_theta_phi
+from popsynth.auxiliary_sampler import NonObservedAuxSampler
+from popsynth.utils.spherical_geometry import sample_theta_phi
 
-# class SkySampler(object):
-#     _auxiliary_sampler_name = "SkySampler"
 
-#     def __init__(self, ra_sampler=None, dec_sampler=None):
-#         """
-#         A sky sampler that samples theta phi positions
+class SkySampler(object):
+    _auxiliary_sampler_name = "SkySampler"
 
-#         :returns:
-#         :rtype:
+    def __init__(
+        self,
+        ra_sampler: NonObservedAuxSampler = None,
+        dec_sampler: NonObservedAuxSampler = None,
+    ):
+        """
+        A sky sampler that samples angular positions
+        in ra and dec. If no samplers are provided, then
+        loads default samplers that sample uniformly on
+        the unit sphere. RA and dec are in radians.
 
-#         """
+        :param ra_sampler: Right ascension (RA) sampler
+        :type ra_sampler: :class:`NonObservedAuxSampler`
+        :param dec_sampler: Declination (Dec) sampler
+        :type dec_sampler: :class:`NonObservedAuxSampler`
+        """
 
-#         self._ra_sampler = ra_sampler
-#         self._dec_sampler = dec_sampler
+        self._ra_sampler = ra_sampler
+        self._dec_sampler = dec_sampler
 
-#         self._setup_sky()
+        self._setup_sky()
 
-#     def _setup_sky(self):
+    def _setup_sky(self):
 
-#         if self._ra_sampler is None:
-#             self._ra_sampler = RASampler()
+        if self._ra_sampler is None:
+            self._ra_sampler = RASampler()
 
-#         if self._dec_sampler is None:
-#             self._dec_sampler = DecSampler()
+        if self._dec_sampler is None:
+            self._dec_sampler = DecSampler()
 
-#     @property
-#     def ra_sampler(self):
+    @property
+    def ra_sampler(self):
 
-#         return self._ra_sampler
+        return self._ra_sampler
 
-#     @property
-#     def dec_sampler(self):
+    @property
+    def dec_sampler(self):
 
-#         return self._dec_sampler
+        return self._dec_sampler
 
-# class RASampler(NonObservedAuxSampler):
-#     _auxiliary_sampler_name = "RASampler"
 
-#     def __init__(self):
-#         """
+class RASampler(NonObservedAuxSampler):
+    _auxiliary_sampler_name = "RASampler"
 
-#         Samples the RA of sky uniformly
+    def __init__(self):
+        """
+        Samples the right ascension (RA)
+        uniformly on the unit sphere.
 
-#         :param true_values:
-#         :returns:
-#         :rtype:
+        RA is in radians.
+        """
 
-#         """
+        super(RASampler, self).__init__(name="ra")
 
-#         super(RASampler, self).__init__(name="ra")
+    def true_sampler(self, size):
 
-#     def true_sampler(self, size):
+        self._true_values = np.random.uniform(0, 2 * np.pi, size=size)
 
-#         self._true_values = np.random.uniform(0, 2 * np.pi, size=size)
 
-# class DecSampler(NonObservedAuxSampler):
-#     _auxiliary_sampler_name = "DecSampler"
+class DecSampler(NonObservedAuxSampler):
+    _auxiliary_sampler_name = "DecSampler"
 
-#     def __init__(self):
-#         """
-#         sampler the dec of the sky uniformly
+    def __init__(self):
+        """
+        Samples the declination (Dec)
+        uniformly on the unit sphere.
 
-#         :param true_values:
-#         :returns:
-#         :rtype:
+        Dec is in radians.
+        """
 
-#         """
+        super(DecSampler, self).__init__(name="dec")
 
-#         super(DecSampler, self).__init__(name="dec")
+    def true_sampler(self, size):
 
-#     def true_sampler(self, size):
-
-#         self._true_values = np.arccos(1 - 2 *
-#                                       np.random.uniform(0.0, 1.0, size=size))
+        self._true_values = np.arccos(1 - 2 *
+                                      np.random.uniform(0.0, 1.0, size=size))
diff --git a/popsynth/aux_samplers/trunc_normal_aux_sampler.py b/popsynth/aux_samplers/trunc_normal_aux_sampler.py
index ab6be23..eb2591a 100644
--- a/popsynth/aux_samplers/trunc_normal_aux_sampler.py
+++ b/popsynth/aux_samplers/trunc_normal_aux_sampler.py
@@ -13,18 +13,30 @@ class TruncatedNormalAuxSampler(AuxiliarySampler):
     upper = AuxiliaryParameter()
     sigma = AuxiliaryParameter(default=1, vmin=0)
 
-    def __init__(self, name: str, observed: bool=True):
+    def __init__(self, name: str, observed: bool = True):
         """
-        A truncated normal distribution
+        A truncated normal sampler,
+        where property ~ N(``mu``, ``sigma``), between
+        ``lower`` and ``upper``.
 
-        :param name: 
+        :param name: Name of the property
         :type name: str
-        :param observed: 
+        :param observed: `True` if the property is observed,
+        `False` if it is latent. Defaults to `True`
         :type observed: bool
-        :returns: 
-
+        :param mu: Mean of the normal
+        :type mu: :class:`AuxiliaryParameter`
+        :param tau: Standard deviation of the normal
+        :type tau: :class:`AuxiliaryParameter`
+        :param lower: Lower bound of the truncation
+        :type lower: :class:`AuxiliaryParameter`
+        :param upper: Upper bound of the truncation
+        :type upper: :class:`AuxiliaryParameter`
+        :param sigma: Standard deviation of normal distribution
+        from which observed values are sampled, if ``observed``
+        is `True`
+        :type sigma: :class:`AuxiliaryParameter`
         """
-        
 
         super(TruncatedNormalAuxSampler, self).__init__(name=name,
                                                         observed=observed)
diff --git a/popsynth/aux_samplers/viewing_angle_sampler.py b/popsynth/aux_samplers/viewing_angle_sampler.py
index 0c767a7..6bc7f77 100644
--- a/popsynth/aux_samplers/viewing_angle_sampler.py
+++ b/popsynth/aux_samplers/viewing_angle_sampler.py
@@ -1,7 +1,6 @@
 import numpy as np
 
-from popsynth.auxiliary_sampler import (AuxiliaryParameter,
-                                        NonObservedAuxSampler)
+from popsynth.auxiliary_sampler import AuxiliaryParameter, NonObservedAuxSampler
 
 
 class ViewingAngleSampler(NonObservedAuxSampler):
@@ -11,13 +10,13 @@ class ViewingAngleSampler(NonObservedAuxSampler):
 
     def __init__(self):
         """
-        A viewing angle sampler that samples from 0, max_angle.
-        It assumes that this is NOT an observed property
-
-        :param max_angle: the maximum angle to which to sample in DEGS
-        :returns: None
-        :rtype: None
+        A viewing angle sampler that samples
+        from 0 to ``max_angle``. Unlike other samplers,
+        it assumes that this is NOT an observed property
 
+        :param max_angle: The maximum angle to which to
+        sample in degrees
+        :type max_angle: :class:`AuxiliaryParameter`
         """
 
         super(ViewingAngleSampler, self).__init__(name="va", )
diff --git a/popsynth/auxiliary_sampler.py b/popsynth/auxiliary_sampler.py
index 7bf66d4..aae3ce3 100644
--- a/popsynth/auxiliary_sampler.py
+++ b/popsynth/auxiliary_sampler.py
@@ -21,15 +21,34 @@ class AuxiliaryParameter(Parameter):
     pass
 
 
-class AuxiliarySampler(object,
-                       metaclass=AutoRegister(auxiliary_parameter_registry,
-                                              base_type=ParameterMeta)):
-    def __init__(self,
-                 name: str,
-                 observed: bool = True,
-                 uses_distance: bool = False,
-                 uses_luminosity: bool = False,
-                 uses_sky_position: bool = False) -> None:
+class AuxiliarySampler(
+        object,
+        metaclass=AutoRegister(auxiliary_parameter_registry,
+                               base_type=ParameterMeta),
+):
+    def __init__(
+        self,
+        name: str,
+        observed: bool = True,
+        uses_distance: bool = False,
+        uses_luminosity: bool = False,
+        uses_sky_position: bool = False,
+    ) -> None:
+        """
+        Base class for auxiliary samplers.
+
+        :param name: Name of the sampler
+        :type name: str
+        :param observed: `True` if the property is observed,
+        `False` if it is latent. Defaults to `True`
+        :type observed: bool
+        :param uses_distance: `True` if sampler uses distance values
+        :type uses_distance: bool
+        :param uses_luminosity: `True` if sampler uses luminosities
+        :type uses_luminosity: bool
+        :param uses_sky_position: `True` if sampler uses sky positions
+        :type uses_sky_position: bool
+        """
 
         self._parameter_storage = {}  # type: Dict[str, float]
         self._name = name  # type: str
@@ -49,23 +68,21 @@ def __init__(self,
         self._uses_sky_position = uses_sky_position  # type: bool
 
     def set_luminosity(self, luminosity: ArrayLike) -> None:
-        """FIXME! briefly describe function
-
-        :param luminosity:
-        :returns:
-        :rtype:
+        """
+        Set the luminosity values.
 
+        :param luminosity: Luminosity
+        :type luminosity: ArrayLike
         """
 
         self._luminosity = luminosity  # type:ArrayLike
 
     def set_spatial_values(self, value: SpatialContainer) -> None:
-        """FIXME! briefly describe function
-
-        :param distance:
-        :returns:
-        :rtype:
+        """
+        Set the spatial values.
 
+        :param value: Spatial values
+        :type value: :class:`SpatialContainer`
         """
 
         self._distance = value.distance  # type:ArrayLike
@@ -85,14 +102,14 @@ def set_selection_probability(self, selector: SelectionProbabilty) -> None:
 
     def _apply_selection(self) -> None:
         """
-        Default selection if none is specfied in child class
+        Default selection if none is specfied in child class.
         """
 
         self._selector.draw(len(self._obs_values))
 
     def set_secondary_sampler(self, sampler) -> None:
         """
-        Allows the setting of a secondary sampler from which to derive values
+        Allows the setting of a secondary sampler from which to derive values.
         """
 
         # make sure we set the sampler as a secondary
@@ -109,7 +126,8 @@ def draw(self, size: int = 1):
         """
         Draw the primary and secondary samplers. This is the main call.
 
-        :param size: the number of samples to draw
+        :param size: The number of samples to draw
+        :type size: int
         """
         # do not resample!
         if not self._is_sampled:
@@ -163,7 +181,7 @@ def draw(self, size: int = 1):
 
             else:
 
-                self._obs_values = (self._true_values)  # type: ArrayLike
+                self._obs_values = self._true_values  # type: ArrayLike
 
             self._selector.set_observed_value(self._obs_values)
 
@@ -185,7 +203,7 @@ def draw(self, size: int = 1):
 
     def reset(self):
         """
-        reset all the selections
+        Reset all the selections.
         """
 
         if self._is_sampled:
@@ -219,12 +237,15 @@ def get_secondary_properties(
         primary=None,
         spatial_distribution=None,
     ) -> SamplerDict:
-        """FIXME! briefly describe function
-
-        :param recursive_secondaries:
-        :returns:
-        :rtype:
-
+        """
+        Get properties of secondary samplers.
+
+        :param recursive_secondaries: Recursive dict of secondaries
+        :param graph: Graph
+        :param primary: Primary sampler
+        :param spatial_distribution: Spatial Distribution
+        :returns: Dict of samplers
+        :rtype: :class:`SamplerDict`
         """
 
         # if a holder was not passed, create one
@@ -267,12 +288,11 @@ def get_secondary_objects(
         self,
         recursive_secondaries: Optional[Dict[str, Any]] = None,
     ) -> Dict[str, Any]:
-        """FIXME! briefly describe function
-
-        :param recursive_secondaries:
-        :returns:
-        :rtype:
+        """Get secondary objects.
 
+        :param recursive_secondaries: Recursive dict of secondaries
+        :returns: Dict of objects
+        :rtype: Dict[str, Any]
         """
 
         # if a holder was not passed, create one
@@ -312,7 +332,9 @@ def get_secondary_objects(
     @property
     def secondary_samplers(self) -> SamplerDict:
         """
-        Secondary samplers
+        Secondary samplers.
+        :returns: Dict of secondary samplers
+        :rtype: :class:`SamplerDict`
         """
 
         return self._secondary_samplers
@@ -333,7 +355,6 @@ def has_secondary(self) -> bool:
 
     @property
     def observed(self) -> bool:
-        """"""
         return self._is_observed
 
     @property
@@ -346,37 +367,14 @@ def obs_name(self) -> str:
 
     @property
     def true_values(self) -> np.ndarray:
-        """
-        The true values
-
-        :returns:
-        :rtype:
-
-        """
-
         return self._true_values
 
     @property
     def obs_values(self) -> np.ndarray:
-        """
-        The observed values
-        :returns:
-        :rtype:
-
-        """
-
         return self._obs_values
 
     @property
     def selection(self) -> np.ndarray:
-        """
-        The selection function
-
-        :returns:
-        :rtype: np.ndarray
-
-        """
-
         return self._selector.selection
 
     @property
@@ -411,7 +409,7 @@ def uses_luminosity(self) -> np.ndarray:
     @property
     def luminosity_distance(self):
         """
-        luminosity distance if needed
+        luminosity distance if needed.
         """
 
         return cosmology.luminosity_distance(self._distance)
@@ -442,14 +440,13 @@ def __init__(self,
 
 class DerivedLumAuxSampler(AuxiliarySampler):
     def __init__(self, name: str, uses_distance: bool = False):
-        """FIXME! briefly describe function
-
-        :param name:
-        :param sigma:
-        :param observed:
-        :returns:
-        :rtype:
+        """
+        Base class for generating luminosity from other properties.
 
+        :param name: Name of the sampler
+        :type name: str
+        :param uses_distance: `True` if sampler uses distance values
+        :type uses_distance: bool
         """
 
         super(DerivedLumAuxSampler, self).__init__(name,
diff --git a/popsynth/distribution.py b/popsynth/distribution.py
index b696fa8..9ee0b46 100644
--- a/popsynth/distribution.py
+++ b/popsynth/distribution.py
@@ -1,6 +1,6 @@
 import abc
 from dataclasses import dataclass
-from typing import Dict, List, Union
+from typing import Dict, Union
 
 import numpy as np
 from class_registry import AutoRegister
@@ -28,15 +28,14 @@ class Distribution(object,
 
     def __init__(self, name: str, seed: int, form: str) -> None:
         """
-        A distribution base class
-
-        :param name: the name of the distribution
-        :param seed: the random seed
-        :param form: the latex form
-        :param truth: dictionary holding true parameters
-        :returns:
-        :rtype:
-
+        A distribution base class.
+
+        :param name: Name of the distribution
+        :type name: str
+        :param seed: Random seed
+        :type seed: int
+        :param form: the LaTeX form
+        :type form: str
         """
 
         self._parameter_storage: Dict[str, float] = {}
@@ -73,6 +72,9 @@ def truth(self) -> Dict[str, float]:
 
 @dataclass
 class SpatialContainer:
+    """
+    Container for 3D spatial values.
+    """
 
     distance: ArrayLike
     theta: ArrayLike
@@ -94,14 +96,15 @@ class SpatialDistribution(Distribution):
 
     def __init__(self, name: str, seed: int, form: Union[str, None] = None):
         """
-        A spatial distribution such as a redshift
-        distribution
-
-        :param name: the name of the distribution
-        :param r_max: the maximum distance to sample
-        :param seed: the random seed
-        :param form: the latex form
-
+        A base class for spatial distributions,
+        such as redshift distributions.
+
+        :param name: Name of the distribution
+        :type name: str
+        :param seed: Random seed
+        :type seed: int
+        :param form: the LaTeX form
+        :type form: Union[str, None]
         """
         self._theta = None
         self._phi = None
@@ -113,12 +116,9 @@ def __init__(self, name: str, seed: int, form: Union[str, None] = None):
     @abc.abstractmethod
     def differential_volume(self, distance):
         """
-        the differential volume
-
-        :param distance:
-        :returns:
-        :rtype:
+        The differential volume
 
+        :param distance: Distance
         """
 
         raise RuntimeError("Must be implemented in derived class")
@@ -129,14 +129,11 @@ def dNdV(self, distance):
         raise RuntimeError("Must be implemented in derived class")
         pass
 
-    def time_adjustment(self, r):
+    def time_adjustment(self, distance):
         """
-        the time adjustment
-
-        :param r:
-        :returns:
-        :rtype:
+        The time adjustment
 
+        :param distance: Distance
         """
 
         return 1.0
@@ -176,7 +173,10 @@ def draw_sky_positions(self, size: int) -> None:
 
     def draw_distance(self, size: int) -> None:
         """
-        Draw the distances from the specified dN/dr model
+        Draw the distances from the specified dN/dr model.
+
+        :param size: Number of distances to sample
+        :type size: int
         """
 
         # create a callback for the sampler
@@ -225,12 +225,14 @@ class LuminosityDistribution(Distribution):
 
     def __init__(self, name: str, seed: int, form: Union[str, None] = None):
         """
-        A luminosity distribution such as a
-        distribution
-
-        :param name: the name of the distribution
-        :param seed: the random seed
-        :param form: the latex form
+        A base class for luminosity distributions.
+
+        :param name: Name of the distribution
+        :type name: str
+        :param seed: Random seed
+        :type seed: int
+        :param form: the LaTeX form
+        :type form: Union[str, None]
         """
 
         super(LuminosityDistribution, self).__init__(
@@ -242,12 +244,9 @@ def __init__(self, name: str, seed: int, form: Union[str, None] = None):
     @abc.abstractmethod
     def phi(self, luminosity):
         """
-        The functional form of the distribution
-
-        :param L:
-        :returns:
-        :rtype:
+        The functional form of the distribution.
 
+        :param luminosity: Luminosity
         """
 
         raise RuntimeError("Must be implemented in derived class")
diff --git a/popsynth/distributions/__init__.py b/popsynth/distributions/__init__.py
index 2d1cbcc..206fdbd 100644
--- a/popsynth/distributions/__init__.py
+++ b/popsynth/distributions/__init__.py
@@ -1,19 +1,36 @@
 from .bpl_distribution import BPLDistribution
-from .cosmological_distribution import (CosmologicalDistribution,
-                                        SFRDistribution,
-                                        ZPowerCosmoDistribution)
+from .cosmological_distribution import (
+    CosmologicalDistribution,
+    SFRDistribution,
+    ZPowerCosmoDistribution,
+)
+from .delta_distribution import DeltaDistribution
+from .flatland_distribution import FlatlandDistribution
 from .log10_normal_distribution import Log10NormalDistribution
 from .log_normal_distribution import LogNormalDistribution
 from .pareto_distribution import ParetoDistribution
 from .schechter_distribution import SchechterDistribution
-from .spherical_distribution import (ConstantSphericalDistribution,
-                                     SphericalDistribution,
-                                     ZPowerSphericalDistribution)
+from .spherical_distribution import (
+    ConstantSphericalDistribution,
+    SphericalDistribution,
+    ZPowerSphericalDistribution,
+)
+from .spiral_galaxy_distribution import SpiralGalaxyDistribution
 
 __all__ = [
-    "SphericalDistribution", "CosmologicalDistribution", "SFRDistribution",
-    "ZPowerCosmoDistribution", "ParetoDistribution", "Log10NormalDistribution",
-    "LogNormalDistribution", "SchechterDistribution", "BPLDistribution",
-    "SphericalDistribution", "ConstantSphericalDistribution",
-    "ZPowerSphericalDistribution"
+    "SphericalDistribution",
+    "CosmologicalDistribution",
+    "SFRDistribution",
+    "ZPowerCosmoDistribution",
+    "ParetoDistribution",
+    "Log10NormalDistribution",
+    "LogNormalDistribution",
+    "SchechterDistribution",
+    "BPLDistribution",
+    "SphericalDistribution",
+    "ConstantSphericalDistribution",
+    "ZPowerSphericalDistribution",
+    "DeltaDistribution",
+    "FlatlandDistribution",
+    "SpiralGalaxyDistribution",
 ]
diff --git a/popsynth/distributions/bpl_distribution.py b/popsynth/distributions/bpl_distribution.py
index 13ee546..9accb68 100644
--- a/popsynth/distributions/bpl_distribution.py
+++ b/popsynth/distributions/bpl_distribution.py
@@ -12,20 +12,31 @@ class BPLDistribution(LuminosityDistribution):
     beta = DistributionParameter()
     Lmax = DistributionParameter(vmin=0)
 
-    def __init__(self, seed=1234, name="bpl"):
-
-        """
-
-        a broken power law distribution 
-
-        :param seed: 
-        :type seed: 
-        :param name: 
-        :type name: 
-        :returns: 
-
+    def __init__(self, seed: int = 1234, name: str = "bpl"):
+        """A broken power law luminosity distribution.
+
+        L ~ L^``alpha`` for L <= ``Lbreak``
+        L ~ L^``beta`` for L > ``Lbreak``
+
+        :param seed: Random seed
+        :type seed: int
+        :param name: Name of the distribution
+        :type name: str
+        :param Lmin: Minimum value of the luminosity
+        :type Lmin: :class:`DistributionParameter`
+        :param alpha: Index of the lower power law
+        :type alpha: :class:`DistributionParameter`
+        :param Lbreak: Luminosity of the power law break
+        :type Lbreak: :class:`DistributionParameter`
+        :param beta: Index of the upper power law
+        :type beta: :class:`DistributionParameter`
+        :param Lmax: Maximum value of the luminosity
+        :type Lmax: :class:`DistributionParameter`
         """
-        lf_form = r"\frac{\alpha L_{\rm min}^{\alpha}}{L^{\alpha+1}}"
+        lf_form = r"\begin{cases} C L^{\alpha} & \mbox{if } L"
+        lf_form += r"\leq L_\mathrm{break},\\ C L^{\beta} "
+        lf_form += r"L_\mathrm{break}^{\alpha - \beta}"
+        lf_form += r" & \mbox{if } L > L_\mathrm{break}. \end{cases}"
 
         super(BPLDistribution, self).__init__(seed=seed,
                                               name=name,
@@ -45,12 +56,13 @@ def draw_luminosity(self, size=1):
 
 def integrate_pl(x0, x1, x2, a1, a2):
     """
-    x0: lower bound
-    x1: break point
-    x2: upper bound
-    a1: lower power law index
-    a2: upper power low index
+    Integrate a broken power law between bounds.
 
+    :param x0: Lower bound
+    :param x1: Break point
+    :param x2: Upper bound
+    :param a1: Lower power law index
+    :param a2: Upper power low index
     """
 
     # compute the integral of each piece analytically
@@ -70,13 +82,14 @@ def integrate_pl(x0, x1, x2, a1, a2):
 
 def bpl(x, x0, x1, x2, a1, a2):
     """
-    x: the domain of the function
-    x0: lower bound
-    x1: break point
-    x2: upper bound
-    a1: lower power law index
-    a2: upper power low index
-
+    Broken power law between bounds.
+
+    :param x: The domain of the function
+    :param x0: Lower bound
+    :param x1: Break point
+    :param x2: Upper bound
+    :param a1: Lower power law index
+    :param a2: Upper power low index
     """
 
     # creatre a holder for the values
@@ -100,12 +113,15 @@ def bpl(x, x0, x1, x2, a1, a2):
 
 def sample_bpl(u, x0, x1, x2, a1, a2):
     """
-    u: uniform random number between on {0,1}
-    x0: lower bound
-    x1: break point
-    x2: upper bound
-    a1: lower power law index
-    a2: upper power low index
+    Sample from a broken power law
+    between bounds.
+
+    :param u: Uniform random number on {0,1}
+    :param x0: Lower bound
+    :param x1: Break point
+    :param x2: Upper bound
+    :param a1: Lower power law index
+    :param a2: Upper power low index
     """
 
     # compute the weights with our integral function
diff --git a/popsynth/distributions/cosmological_distribution.py b/popsynth/distributions/cosmological_distribution.py
index 670879e..8683cdb 100644
--- a/popsynth/distributions/cosmological_distribution.py
+++ b/popsynth/distributions/cosmological_distribution.py
@@ -1,7 +1,6 @@
 import numpy as np
 import numba as nb
-import math
-
+from typing import Dict, Any
 from popsynth.distribution import SpatialDistribution, DistributionParameter
 from popsynth.utils.cosmology import cosmology
 
@@ -9,28 +8,30 @@
 class CosmologicalDistribution(SpatialDistribution):
     _distribution_name = "CosmologicalDistribution"
 
-    def __init__(self,
-                 seed=1234,
-                 name="cosmo",
-                 form=None,
-                 truth={},
-                 is_rate=True):
-
+    def __init__(
+        self,
+        seed: int = 1234,
+        name: str = "cosmo",
+        form: str = None,
+        truth: Dict[str, Any] = {},
+        is_rate: bool = True,
+    ):
         """
-        base cosmological spatial distribution
-
-        :param seed: 
-        :type seed: 
-        :param name: 
-        :type name: 
-        :param form: 
-        :type form: 
-        :param truth: 
-        :type truth: 
-        :param is_rate: 
-        :type is_rate: 
-        :returns: 
-
+        Base class for cosmological spatial distributions.
+
+        :param seed: Random seed
+        :type seed: int
+        :param name: Name of the distribution
+        :type name: str
+        :param form: Mathematical description of distribution
+        :type form: str
+        :param truth: True values of parameters
+        :type truth: dict[str, Any]
+        :param is_rate: `True` if modelling a population of transient events,
+        `False` if modelling a population of steady-state objects.
+        Affects the ``time_adjustment`` method used in cosmo calculations.
+        Default is `True`.
+        :type is_rate: bool
         """
         super(CosmologicalDistribution, self).__init__(
             seed=seed,
@@ -40,17 +41,48 @@ def __init__(self,
         self._is_rate = is_rate
 
     def differential_volume(self, z):
+        """
+        Differential comoving volume in Gpc^3 sr^-1.
+
+        dV/dzdOmega
+
+        :param z: Redshift
+        :returns: The differential comoving volume in
+        Gpc^-3 sr^-1.
+        """
 
         return cosmology.differential_comoving_volume(z)
 
     def transform(self, L, z):
+        """
+        Transformation from luminosity to energy flux.
+
+        L / 4 pi dL^2
+
+        dL is in cm. Therefore for L in erg s^-1 returns
+        flux in erg cm^-2 s^-1.
+
+        :param L: Luminosity
+        :param z: Redshift
+        :returns: Flux
+        """
 
         return L / (4.0 * np.pi * cosmology.luminosity_distance(z)**2)
 
     def time_adjustment(self, z):
+        """
+        Time adjustment factor to handle both
+        transient and steady-state populations.
+
+        :param z: Redshift
+        :returns: Appropriate factor depending on ``is_rate``
+        """
         if self._is_rate:
+
             return 1 + z
+
         else:
+
             return 1.0
 
 
@@ -63,22 +95,37 @@ class SFRDistribution(CosmologicalDistribution):
     decay = DistributionParameter()
     peak = DistributionParameter(vmin=0)
 
-    def __init__(self, seed=1234, name="sfr", is_rate=True):
-
+    def __init__(self,
+                 seed: int = 1234,
+                 name: str = "sfr",
+                 is_rate: bool = True):
         """
         A star-formation like distribution of the form
-        presented in Cole et al. 2001
-
-        :param seed: 
-        :type seed: 
-        :param name: 
-        :type name: 
-        :param is_rate: 
-        :type is_rate: 
-        :returns: 
-
+        presented in Cole et al. 2001.
+
+        ``r0``(``a``+``rise``z)/(1 + (z/``peak``)^``decay``)
+
+        :param seed: Random seed
+        :type seed: int
+        :param name: Name of the distribution
+        :type name: str
+        :param is_rate: `True` if modelling a population of transient events,
+        `False` if modelling a population of steady-state objects.
+        Affects the ``time_adjustment`` method used in cosmo calculations.
+        Default is `True`.
+        :type is_rate: bool
+        :param r0: The local density in units of Gpc^-3
+        :type r0: :class:`DistributionParameter`
+        :param a: Offset at z=0
+        :type a: :class:`DistributionParameter`
+        :param rise: Rise at low z
+        :type rise: :class:`DistributionParameter`
+        :param decay: Decay at high z
+        :type decay: :class:`DistributionParameter`
+        :param peak: Peak of z distribution
+        :type peak: :class:`DistributionParameter`
         """
-        spatial_form = r"\rho_0 \frac{1+r \cdot z}{1+ \left(z/p\right)^d}"
+        spatial_form = r"\rho_0 \frac{a+r \cdot z}{1+ \left(z/p\right)^d}"
 
         super(SFRDistribution, self).__init__(
             seed=seed,
@@ -88,7 +135,7 @@ def __init__(self, seed=1234, name="sfr", is_rate=True):
         )
 
     def dNdV(self, z):
-        return _dndv(
+        return _sfr_dndv(
             z,
             self.r0,
             self.a,
@@ -99,7 +146,7 @@ def dNdV(self, z):
 
 
 @nb.njit(fastmath=True)
-def _dndv(z, r0, a, rise, decay, peak):
+def _sfr_dndv(z, r0, a, rise, decay, peak):
     top = a + rise * z
     bottom = 1.0 + np.power(z / peak, decay)
 
@@ -114,23 +161,29 @@ class ZPowerCosmoDistribution(CosmologicalDistribution):
 
     def __init__(
         self,
-        seed=1234,
-        name="zpow_cosmo",
-        is_rate=True,
+        seed: int = 1234,
+        name: str = "zpow_cosmo",
+        is_rate: bool = True,
     ):
-
         """
-        a cosmological distribution where the density 
-        evolves as a power law
-
-        :param seed: 
-        :type seed: 
-        :param name: 
-        :type name: 
-        :param is_rate: 
-        :type is_rate: 
-        :returns: 
-
+        A cosmological distribution where the density
+        evolves as a power law.
+
+        ``Lambda`` (1+z)^``delta``
+
+        :param seed: Random seed
+        :type seed: int
+        :param name: Name of the distribution
+        :type name: str
+        :param is_rate: `True` if modelling a population of transient events,
+        `False` if modelling a population of steady-state objects.
+        Affects the ``time_adjustment`` method used in cosmo calculations.
+        Default is `True`.
+        :type is_rate: bool
+        :param Lambda: The local density in units of Gpc^-3
+        :type Lambda: :class:`DistributionParameter`
+        :param delta: The index of the power law
+        :type delta: :class:`DistributionParameter`
         """
         spatial_form = r"\Lambda (z+1)^{\delta}"
 
@@ -152,7 +205,7 @@ def _zp_dndv(z, Lambda, delta):
 
 
 # class MergerDistribution(CosmologicalDistribution):
-#_distribution_name = "MergerDistribution"
+# _distribution_name = "MergerDistribution"
 #     def __init__(self, r0, td, sigma, r_max=10, seed=1234, name="merger"):
 
 #         spatial_form = r"\rho_0 \frac{1+r \cdot z}{1+ \left(z/p\right)^d}"
diff --git a/popsynth/distributions/delta_distribution.py b/popsynth/distributions/delta_distribution.py
index f1077ba..a04ad5d 100644
--- a/popsynth/distributions/delta_distribution.py
+++ b/popsynth/distributions/delta_distribution.py
@@ -8,17 +8,17 @@ class DeltaDistribution(LuminosityDistribution):
 
     Lp = DistributionParameter(vmin=0, default=0)
 
-    def __init__(self, seed=1234, name="delta"):
-
+    def __init__(self, seed: int = 1234, name: str = "delta"):
         """
-        a dirac delta distribution
-
-        :param seed: 
-        :type seed: 
-        :param name: 
-        :type name: 
-        :returns: 
-
+        A delta function luminosity distribution,
+        centred on ``Lp``.
+
+        :param seed: Random seed
+        :type seed: int
+        :param name: Name of the distribution
+        :type name: str
+        :param Lp: The central value
+        :type Lp: :class:`DistributionParameter`
         """
         lf_form = r"\delta(L - L_p)"
 
diff --git a/popsynth/distributions/flatland_distribution.py b/popsynth/distributions/flatland_distribution.py
index 2285e61..8e5836f 100644
--- a/popsynth/distributions/flatland_distribution.py
+++ b/popsynth/distributions/flatland_distribution.py
@@ -1,5 +1,3 @@
-import numpy as np
-
 from popsynth.distribution import SpatialDistribution, DistributionParameter
 
 
@@ -8,19 +6,21 @@ class FlatlandDistribution(SpatialDistribution):
 
     Lambda = DistributionParameter(default=1, vmin=0)
 
-    def __init__(self, seed: int = 1234, name: str = "flatland", form=None):
-
+    def __init__(self,
+                 seed: int = 1234,
+                 name: str = "flatland",
+                 form: str = None):
         """
-        a distribution with only length
+        A flat spatial distribution with only length.
 
-        :param seed: 
+        :param seed: Random seed
         :type seed: int
-        :param name: 
+        :param name: Name of the distribution
         :type name: str
-        :param form: 
-        :type form: 
-        :returns: 
-
+        :param form: Mathematical description of distribution
+        :type form: str
+        :param Lambda: Length
+        :type Lambda: :class:`DistributionParameter`
         """
         super(FlatlandDistribution, self).__init__(seed=seed,
                                                    name=name,
diff --git a/popsynth/distributions/log10_normal_distribution.py b/popsynth/distributions/log10_normal_distribution.py
index d3ae379..2510cf5 100644
--- a/popsynth/distributions/log10_normal_distribution.py
+++ b/popsynth/distributions/log10_normal_distribution.py
@@ -9,19 +9,22 @@ class Log10NormalDistribution(LuminosityDistribution):
     mu = DistributionParameter()
     tau = DistributionParameter(vmin=0)
 
-    def __init__(self, seed=1234, name="log10norm"):
-
+    def __init__(self, seed: int = 1234, name: str = "log10norm"):
         """
-        a log10-normal luminosity function
-
-        :param seed: 
-        :type seed: 
-        :param name: 
-        :type name: 
-        :returns: 
-
+        A log10-normal luminosity function
+
+        Log10Normal(``mu``, ``tau``)
+
+        :param seed: Random seed
+        :type seed: int
+        :param name: Name of the distribution
+        :type name: str
+        :param mu: Mean of the log10 normal
+        :type mu: :class:`DistributionParameter`
+        :param tau: Standard deviation of the log10 normal
+        :type tau: :class:`DistributionParameter`
         """
-        lf_form = r"\frac{\alpha L_{\rm min}^{\alpha}}{L^{\alpha+1}}"
+        lf_form = r"\mathrm{Log_{10}Normal(\mu, \tau)}"
 
         super(Log10NormalDistribution, self).__init__(name=name,
                                                       seed=seed,
diff --git a/popsynth/distributions/log_normal_distribution.py b/popsynth/distributions/log_normal_distribution.py
index a42676e..00ca0e0 100644
--- a/popsynth/distributions/log_normal_distribution.py
+++ b/popsynth/distributions/log_normal_distribution.py
@@ -1,5 +1,4 @@
 import numpy as np
-import scipy.stats as stats
 
 from popsynth.distribution import LuminosityDistribution, DistributionParameter
 
@@ -10,23 +9,28 @@ class LogNormalDistribution(LuminosityDistribution):
     mu = DistributionParameter()
     tau = DistributionParameter(vmin=0)
 
-    def __init__(self, seed=1234, name="lognorm"):
-
+    def __init__(self, seed: int = 1234, name: str = "lognorm"):
         """
-        a log-normal luminosity distribution
-
-        :param seed: 
-        :type seed: 
-        :param name: 
-        :type name: 
-        :returns: 
-
+        A log-normal luminosity distribution.
+
+        LogNormal(``mu``, ``tau``)
+
+        :param seed: Random seed
+        :type seed: int
+        :param name: Name of the distribution
+        :type name: str
+        :param mu: Mean of the log normal
+        :type mu: :class:`DistributionParameter`
+        :param tau: Standard deviation of the log normal
+        :type tau: :class:`DistributionParameter`
         """
-        lf_form = r"\frac{\alpha L_{\rm min}^{\alpha}}{L^{\alpha+1}}"
+        lf_form = r"\mathrm{LogNorm}(\mu, \tau)"
 
-        super(LogNormalDistribution, self).__init__(name=name,
-                                                    seed=seed,
-                                                    form=lf_form)
+        super(LogNormalDistribution, self).__init__(
+            name=name,
+            seed=seed,
+            form=lf_form,
+        )
 
     def phi(self, L):
 
diff --git a/popsynth/distributions/pareto_distribution.py b/popsynth/distributions/pareto_distribution.py
index 0a716b8..bde705b 100644
--- a/popsynth/distributions/pareto_distribution.py
+++ b/popsynth/distributions/pareto_distribution.py
@@ -11,14 +11,19 @@ class ParetoDistribution(LuminosityDistribution):
 
     def __init__(self, seed: int = 1234, name: str = "pareto"):
         """
-        A Pareto luminosity function
-        :param seed:
-        :param name:
-        :returns:
-        :rtype:
-
+        A Pareto luminosity function.
+
+        ``alpha``*``Lmin``^``alpha`` / L^(``alpha``+1)
+
+        :param seed: Random seed
+        :type seed: int
+        :param name: Name of the distribution
+        :type name: str
+        :param Lmin: Minimum value of the luminosity
+        :type Lmin: :class:`DistributionParameter`
+        :param alpha: Index of the pareto distribution
+        :type alpha: :class:`DistributionParameter`
         """
-
         lf_form = r"\frac{\alpha L_{\rm min}^{\alpha}}{L^{\alpha+1}}"
 
         super(ParetoDistribution, self).__init__(seed=seed,
@@ -26,14 +31,6 @@ def __init__(self, seed: int = 1234, name: str = "pareto"):
                                                  form=lf_form)
 
     def phi(self, L):
-        """
-        The luminosity function
-
-        :param L:
-        :returns:
-        :rtype:
-
-        """
 
         out = np.zeros_like(L)
 
@@ -45,13 +42,5 @@ def phi(self, L):
         return out
 
     def draw_luminosity(self, size: int = 1):
-        """
-        draw the luminosity
-
-        :param size:
-        :returns:
-        :rtype:
-
-        """
 
         return (np.random.pareto(self.alpha, size) + 1) * self.Lmin
diff --git a/popsynth/distributions/schechter_distribution.py b/popsynth/distributions/schechter_distribution.py
index 14f8fb7..aabbcab 100644
--- a/popsynth/distributions/schechter_distribution.py
+++ b/popsynth/distributions/schechter_distribution.py
@@ -10,44 +10,36 @@ class SchechterDistribution(LuminosityDistribution):
     Lmin = DistributionParameter(default=1, vmin=0)
     alpha = DistributionParameter(default=1)
 
-    def __init__(self, seed=1234, name="schechter"):
+    def __init__(self, seed: int = 1234, name: str = "schechter"):
         """
-        A Schechter luminosity function
-        :param seed: the random number seed
-        :param name: the name
-        :returns: None
-        :rtype:
-
+        A Schechter luminosity function as in
+        Schechter, Astrophysical Journal, Vol. 203,
+        p. 297-306 (1976).
+
+        :param seed: Random seed
+        :type seed: int
+        :param name: Name of the distribution
+        :type name: str
+        :param Lmin: Minimum value of the luminosity
+        :type Lmin: :class:`DistributionParameter`
+        :param alpha: Index of the distribution
+        :type alpha: :class:`DistributionParameter`
         """
+        lf_form = r"\frac{1}{L_{\rm min}^{1+\alpha}\Gamma\left(1+\alpha\right)}"
+        lf_form += r" L^{\alpha} \exp\left[ - \frac{L}{L_{\rm min}}\right]"
 
-        form = r"\frac{1}{L_{\rm min}^{1+\alpha} \Gamma\left(1+\alpha\right)} L^{\alpha} \exp\left[ - \frac{L}{L_{\rm min}}\right]"
-
-        super(SchechterDistribution, self).__init__(name=name,
-                                                    seed=seed,
-                                                    form=form)
+        super(SchechterDistribution, self).__init__(
+            name=name,
+            seed=seed,
+            form=lf_form,
+        )
 
     def phi(self, L):
-        """
-        the functional form
-
-        :param L:
-        :returns:
-        :rtype:
-
-        """
 
         return (L**self.alpha * np.exp(-L / self.Lmin) /
                 (self.Lmin**(1 + self.alpha) * sf.gamma(1 + self.alpha)))
 
     def draw_luminosity(self, size=1):
-        """
-        function to sample the distribution
-
-        :param size:
-        :returns:
-        :rtype:
-
-        """
 
         xs = 1 - np.random.uniform(size=size)
         return self.Lmin * sf.gammaincinv(1 + self.alpha, xs)
diff --git a/popsynth/distributions/spherical_distribution.py b/popsynth/distributions/spherical_distribution.py
index a75fbcf..aab38e8 100644
--- a/popsynth/distributions/spherical_distribution.py
+++ b/popsynth/distributions/spherical_distribution.py
@@ -6,24 +6,28 @@
 class SphericalDistribution(SpatialDistribution):
     _distribution_name = "SphericalDistribution"
 
-    def __init__(self, seed=1234, name="sphere", form=None):
-
+    def __init__(
+        self,
+        seed: int = 1234,
+        name: str = "sphere",
+        form: str = None,
+    ):
         """
-        a generic spherical distribution. Can be inherited to
+        A generic spherical distribution. Can be inherited to
         form more complex spherical distributions
 
-        :param seed: 
-        :type seed: 
-        :param name: 
-        :type name: 
-        :param form: 
-        :type form: 
-        :returns: 
-
+        :param seed: Random seed
+        :type seed: int
+        :param name: Name of the distribution
+        :type name: str
+        :param form: Mathematical description of distribution
+        :type form: str
         """
-        super(SphericalDistribution, self).__init__(seed=seed,
-                                                    name=name,
-                                                    form=form)
+        super(SphericalDistribution, self).__init__(
+            seed=seed,
+            name=name,
+            form=form,
+        )
 
     def differential_volume(self, r):
 
@@ -43,22 +47,21 @@ class ConstantSphericalDistribution(SphericalDistribution):
 
     def __init__(
         self,
-        seed=1234,
-        name="cons_sphere",
-        form=None,
+        seed: int = 1234,
+        name: str = "cons_sphere",
+        form: str = None,
     ):
-
         """
-        A spherical distribution with constant density
-
-        :param seed: 
-        :type seed: 
-        :param name: 
-        :type name: 
-        :param form: 
-        :type form: 
-        :returns: 
-
+        A spherical distribution with constant density.
+
+        :param seed: Random seed
+        :type seed: int
+        :param name: Name of the distribution
+        :type name: str
+        :param form: Mathematical description of distribution
+        :type form: str
+        :param Lambda: Density per unit volume
+        :type Lambda: :class:`DistributionParameter`
         """
         if form is None:
             form = r"\Lambda"
@@ -77,23 +80,28 @@ class ZPowerSphericalDistribution(ConstantSphericalDistribution):
 
     delta = DistributionParameter(default=1)
 
-    def __init__(self, seed=1234, name="zpow_sphere"):
-
+    def __init__(self, seed: int = 1234, name: str = "zpow_sphere"):
         """
-        a spherical population with a power law denisty profile
-
-        :param seed: 
-        :type seed: 
-        :param name: 
-        :type name: 
-        :returns: 
-
+        A spherical distribution with a power law density profile.
+
+        ``Lambda`` (1+r)^``delta``
+
+        :param seed: Random seed
+        :type seed: int
+        :param name: Name of the distribution
+        :type name: str
+        :param form: Mathematical description of distribution
+        :type form: str
+        :param delta: Index of power law distribution
+        :type delta: :class:`DistributionParameter`
         """
-        spatial_form = r"\Lambda (z+1)^{\delta}"
+        spatial_form = r"\Lambda (1+r)^{\delta}"
 
-        super(ZPowerSphericalDistribution, self).__init__(seed,
-                                                          name,
-                                                          form=spatial_form)
+        super(ZPowerSphericalDistribution, self).__init__(
+            seed,
+            name,
+            form=spatial_form,
+        )
 
     def dNdV(self, distance):
 
diff --git a/popsynth/distributions/spiral_galaxy_distribution.py b/popsynth/distributions/spiral_galaxy_distribution.py
index 9649f86..08f8b89 100644
--- a/popsynth/distributions/spiral_galaxy_distribution.py
+++ b/popsynth/distributions/spiral_galaxy_distribution.py
@@ -14,23 +14,41 @@ class SpiralGalaxyDistribution(SphericalDistribution):
     R1 = DistributionParameter()
     R0 = DistributionParameter()
 
-    def __init__(self, seed=1234, name="spiral_galaxy", form=None):
-
+    def __init__(
+        self,
+        seed: int = 1234,
+        name: str = "spiral_galaxy",
+        form: str = None,
+    ):
+        """
+        A spiral galaxy spatial distribution.
+
+        :param seed: Random seed
+        :type seed: int
+        :param name: Name of the distribution
+        :type name: str
+        :param form: Mathematical description of distribution
+        :type form: str
+        :param rho: Local density
+        :type rho: :class:`DistributionParameter`
+        :param a: Shape parameter
+        :type a: :class:`DistributionParameter`
+        :param b: Shape parameter
+        :type b: :class:`DistributionParameter`
+        :param R1: Scale parameter
+        :type R1: :class:`DistributionParameter`
+        :param R0: Scale parameter
+        :type R0: :class:`DistributionParameter`
         """
-        A spiral galaxy spatial distribution
 
-        :param seed: 
-        :type seed: 
-        :param name: 
-        :type name: 
-        :param form: 
-        :type form: 
-        :returns: 
+        spatial_form = r"\rho \Big(\frac{r + R_1}{R_0+R_1}\Big)^a "
+        spatial_form += r"\exp\Big[-b \frac{r-R_0}{R_0 + R_1}\Big]"
 
-        """
-        super(SpiralGalaxyDistribution, self).__init__(seed=seed,
-                                                       name=name,
-                                                       form=form)
+        super(SpiralGalaxyDistribution, self).__init__(
+            seed=seed,
+            name=name,
+            form=form,
+        )
 
     def dNdV(self, r):
 
@@ -40,11 +58,15 @@ def dNdV(self, r):
         ) * np.exp(-self.b * (r - self.R0) / (self.R0 + self.R1)))
 
     def draw_sky_positions(self, size):
+        """
+        Based on Wainscoat 1992 and
+        Faucher-Giguere 2007.
 
-        # code thanks to Moritz!
+        Code thanks to Mortiz Pleintinger.
+        """
 
-        k = np.array([4.25, 4.25, 4.89, 4.89
-                      ])  # Parameter fuer possun [-8.5,0,0], von Wainscoat1992
+        # Parameters for possun [-8.5,0,0], from Wainscoat 1992
+        k = np.array([4.25, 4.25, 4.89, 4.89])
         r0 = np.array([3.48, 3.48, 4.90, 4.90])
         theta0 = np.array([0.0, 3.141, 2.525, 5.666])  # Wainscoat 1992
         height = 0.3  # kpc
@@ -64,8 +86,6 @@ def draw_sky_positions(self, size):
 
         spiraltheta = tet + corrtet  # Faucher-Giguere 2007
 
-        spiralthetaindec = spiraltheta * 180 / np.pi
-
         zpos = rd.exponential(height, size=size)
 
         zpos *= np.random.choice([-1, 1], size=size)
@@ -76,4 +96,5 @@ def draw_sky_positions(self, size):
         phi = np.arccos(zpos / np.sqrt(self._distances**2 + zpos**2))
 
         self._theta = spiraltheta
+
         self._phi = phi
diff --git a/popsynth/population.py b/popsynth/population.py
index 351948a..b9f76ee 100644
--- a/popsynth/population.py
+++ b/popsynth/population.py
@@ -1,9 +1,7 @@
-import importlib
-from typing import Any, Dict, List, Optional, Type
+from typing import Any, Dict, Optional
 
 import h5py
 import ipyvolume as ipv
-import ipywidgets as widgets
 import matplotlib.pyplot as plt
 import networkx as nx
 import numpy as np
@@ -15,9 +13,11 @@
 
 from popsynth.utils.array_to_cmap import array_to_cmap
 from popsynth.utils.hdf5_utils import (
-    clean_graph_dict, fill_graph_dict,
+    clean_graph_dict,
+    fill_graph_dict,
     recursively_load_dict_contents_from_group,
-    recursively_save_dict_contents_to_group)
+    recursively_save_dict_contents_to_group,
+)
 from popsynth.utils.logging import setup_logger
 from popsynth.utils.spherical_geometry import xyz
 
@@ -29,57 +29,84 @@
 
 
 class Population(object):
-    def __init__(self,
-                 luminosities: ArrayLike,
-                 distances: ArrayLike,
-                 known_distances: ArrayLike,
-                 known_distance_idx: ArrayLike,
-                 unknown_distance_idx: ArrayLike,
-                 fluxes: ArrayLike,
-                 flux_obs: ArrayLike,
-                 selection: ArrayLike,
-                 flux_sigma: float,
-                 r_max: float,
-                 n_model: int,
-                 lf_params: Dict[str, Any],
-                 spatial_params: Optional[Dict[str, Any]] = None,
-                 model_spaces: Optional[ArrayLike] = None,
-                 seed: int = 1234,
-                 name: Optional[str] = None,
-                 spatial_form: Optional[Dict[str, Any]] = None,
-                 lf_form: Optional[Dict[str, Any]] = None,
-                 auxiliary_quantities: Optional[Dict[str, Any]] = None,
-                 truth: Dict[str, float] = {},
-                 graph: Optional[Dict[str, Any]] = None,
-                 theta=None,
-                 phi=None,
-                 pop_synth: Optional[Dict[str, Any]] = None) -> None:
+    def __init__(
+        self,
+        luminosities: ArrayLike,
+        distances: ArrayLike,
+        known_distances: ArrayLike,
+        known_distance_idx: ArrayLike,
+        unknown_distance_idx: ArrayLike,
+        fluxes: ArrayLike,
+        flux_obs: ArrayLike,
+        selection: ArrayLike,
+        flux_sigma: float,
+        r_max: float,
+        n_model: int,
+        lf_params: Dict[str, Any],
+        spatial_params: Optional[Dict[str, Any]] = None,
+        model_spaces: Optional[ArrayLike] = None,
+        seed: int = 1234,
+        name: Optional[str] = None,
+        spatial_form: Optional[Dict[str, Any]] = None,
+        lf_form: Optional[Dict[str, Any]] = None,
+        auxiliary_quantities: Optional[Dict[str, Any]] = None,
+        truth: Dict[str, float] = {},
+        graph: Optional[Dict[str, Any]] = None,
+        theta=None,
+        phi=None,
+        pop_synth: Optional[Dict[str, Any]] = None,
+    ) -> None:
         """
         A population container for all the properties of the population.
 
-        :param luminosities: the luminosities
-        :param distances: the distances
-        :param known_distances: the known distances
-        :param known_distance_idx: the index of the known distances
-        :param unknown_distance_idx:  the index of the unknown distances
-        :param fluxes: the latent fluxes
-        :param flux_obs: the observed fluxes
-        :param selection: the selection vector
-        :param flux_sigma: the uncertainty on the observed flux
-        :param r_max: the maximum distance of the survey
-        :param n_model:
-        :param lf_params:
-        :param spatial_params:
-        :param model_spaces:
-        :param seed: the random seed
-        :param name:
-        :param spatial_form:
-        :param lf_form:
-        :param auxiliary_quantities:
-        :param truth:
-        :returns:
-        :rtype:
-
+        :param luminosities: The luminosities
+        :type luminosities: ArrayLike
+        :param distances: The distances
+        :type distances: ArrayLike
+        :param known_distances: The known distances
+        :type known_distances: ArrayLike
+        :param known_distance_idx: The index of the known distances
+        :type known_distance_idx: ArrayLike
+        :param unknown_distance_idx: The index of the unknown distances
+        :type unknown_distance_idx: ArrayLike
+        :param fluxes: The latent fluxes
+        :type fluxes: ArrayLike
+        :param flux_obs: The observed fluxes
+        :type flux_obs: ArrayLike
+        :param selection: The selection vector
+        :type selection: ArrayLike
+        :param flux_sigma: The uncertainty on the observed flux
+        :type flux_sigma: float
+        :param r_max: The maximum distance of the survey
+        :type r_max: float
+        :param n_model: Number of models
+        :type n_model: int
+        :param lf_params: Luminosity function parameters
+        :type lf_params: Dict[str, Any]
+        :param spatial_params: Spatial distribution parameters
+        :type spatial_params: Optional[Dict[str, Any]]
+        :param model_spaces: Model spaces
+        :type model_spaces: ArrayLike
+        :param seed: Random seed
+        :type seed: int
+        :param name: Name of the population
+        :type name: str
+        :param spatial_form: Form of the spatial distribution
+        :type spatial_form: Optional[Dict[str, Any]]
+        :param lf_form: Form of the luminosity function
+        :type lf_form: Optional[Dict[str, Any]]
+        :param auxiliary_quantities: Dict of auxiliary quantities
+        :type auxiliary_quantities: Optional[Dict[str, Any]]
+        :param truth: Dict of true values
+        :type truth: Dict[str, float]
+        :param graph: Graph of generative model
+        :type graph: Optional[Dict[str, Any]]
+        :param theta: Theta
+        :type theta:
+        :param phi: Phi
+        :type phi:
+        :param pop_synth: Population synth
+        :type pop_synth: Optional[Dict[str, Any]]
         """
         self._luminosities = luminosities  # type: ArrayLike
 
@@ -172,11 +199,7 @@ def pop_synth(self) -> Dict[str, Any]:
     @property
     def truth(self):
         """
-        the simulated truth parameters
-
-        :returns:
-        :rtype:
-
+        The simulated truth parameters
         """
 
         return self._truth
@@ -184,51 +207,36 @@ def truth(self):
     @property
     def flux_sigma(self) -> float:
         """
-        the simulated flux sigma
-
-        :returns: 
-
+        The simulated flux sigma
         """
-        
+
         return self._flux_sigma
 
     @property
     def theta(self) -> np.ndarray:
         """
-        the polar angle of the objects
-
-        :returns: 
-
+        The polar angle of the objects
         """
         return self._theta
 
     @property
     def phi(self) -> np.ndarray:
         """
-        the phi angle of the objects
-
-        :returns: 
-
+        The phi angle of the objects
         """
         return self._phi
 
     @property
     def dec(self) -> np.ndarray:
         """
-        the declenation of the objects
-
-        :returns: 
-
+        The declination of the objects
         """
         return 90 - np.rad2deg(self._theta)
 
     @property
     def ra(self) -> np.ndarray:
         """
-        the right ascension of the objects
-
-        :returns: 
-
+        The right ascension of the objects
         """
         return np.rad2deg(self._phi)
 
@@ -236,7 +244,7 @@ def ra(self) -> np.ndarray:
     def luminosities_latent(self) -> np.ndarray:
         """
         The true luminosities of the objects. These are always latent
-        as one cannot directly observe them
+        as one cannot directly observe them.
         """
         return self._luminosities
 
@@ -274,7 +282,6 @@ def fluxes_observed(self) -> np.ndarray:
         """
         All of the observed fluxes, i.e.,
         scattered with error
-
         """
         return self._flux_obs
 
@@ -298,7 +305,7 @@ def selected_fluxes_latent(self) -> np.ndarray:
     def selected_distances(self) -> np.ndarray:
         """
         The selected distances. Note, this is different than
-        the KNOWN distances
+        the KNOWN distances.
         """
         return self._distance_selected
 
@@ -427,12 +434,10 @@ def to_stan_data(self) -> dict:
 
     def writeto(self, file_name: str) -> None:
         """
-        write population to an HDF5 file
-
-        :param file_name:
-        :returns:
-        :rtype:
+        Write population to an HDF5 file
 
+        :param file_name: Name of the file
+        :type file_name: str
         """
 
         with h5py.File(file_name, "w") as f:
@@ -441,8 +446,13 @@ def writeto(self, file_name: str) -> None:
 
     def addto(self, file_name: str, group_name: str) -> None:
         """
-        write population to a group in an existing
+        Write population to a group in an existing
         HDF5 file.
+
+        :param file_name: Name of the file
+        :type file_name: str
+        :param group_name: Name of the group
+        :type group_name: str
         """
 
         with h5py.File(file_name, "r+") as f:
@@ -453,7 +463,7 @@ def addto(self, file_name: str, group_name: str) -> None:
 
     def _writeto(self, f):
         """
-        write to file or group
+        Write to an HDF5 file or group.
 
         :param f: h5py file or group handle
         """
@@ -535,18 +545,15 @@ def _writeto(self, f):
         recursively_save_dict_contents_to_group(
             f, "graph", fill_graph_dict(nx.to_dict_of_dicts(self._graph)))
 
+        # now store the popsynth
         recursively_save_dict_contents_to_group(f, "popsynth", self._pop_synth)
 
-        # now store the popsymnth
-
     @classmethod
-    def from_file(cls, file_name):
+    def from_file(cls, file_name: str):
         """
-        load a population from a file
-        :param file_name:
-        :returns:
-        :rtype:
-
+        Load a population from a file
+        :param file_name: Name of the file
+        :type file_name: str
         """
 
         with h5py.File(file_name, "r") as f:
@@ -554,10 +561,15 @@ def from_file(cls, file_name):
             return cls._loadfrom(f)
 
     @classmethod
-    def from_group(cls, file_name, group_name):
+    def from_group(cls, file_name: str, group_name: str):
         """
-        load a population from a group
+        Load a population from a group
         in a file.
+
+        :param file_name: Name of the file
+        :type file_name: str
+        :param group_name: Name of the group
+        :type group_name: str
         """
 
         with h5py.File(file_name, "r") as f:
@@ -569,7 +581,7 @@ def from_group(cls, file_name, group_name):
     @classmethod
     def _loadfrom(cls, f):
         """
-        load from file or group
+        Load from and HDF5 file or group.
 
         :param f: h5py file or group handle
         """
@@ -652,38 +664,42 @@ def _loadfrom(cls, f):
 
         pop_synth = recursively_load_dict_contents_from_group(f, "popsynth")
 
-        return cls(luminosities=luminosities,
-                   distances=distances,
-                   known_distances=known_distances,
-                   known_distance_idx=known_distance_idx,
-                   unknown_distance_idx=unknown_distance_idx,
-                   fluxes=fluxes,
-                   flux_obs=flux_obs,
-                   selection=selection,
-                   flux_sigma=flux_sigma,
-                   n_model=n_model,
-                   r_max=r_max,
-                   lf_params=lf_params,
-                   spatial_params=spatial_params,
-                   model_spaces=model_spaces,
-                   seed=seed,
-                   name=name,
-                   spatial_form=spatial_form,
-                   lf_form=lf_form,
-                   auxiliary_quantities=auxiliary_quantities,
-                   truth=truth,
-                   graph=graph,
-                   theta=theta,
-                   phi=phi,
-                   pop_synth=pop_synth)
+        return cls(
+            luminosities=luminosities,
+            distances=distances,
+            known_distances=known_distances,
+            known_distance_idx=known_distance_idx,
+            unknown_distance_idx=unknown_distance_idx,
+            fluxes=fluxes,
+            flux_obs=flux_obs,
+            selection=selection,
+            flux_sigma=flux_sigma,
+            n_model=n_model,
+            r_max=r_max,
+            lf_params=lf_params,
+            spatial_params=spatial_params,
+            model_spaces=model_spaces,
+            seed=seed,
+            name=name,
+            spatial_form=spatial_form,
+            lf_form=lf_form,
+            auxiliary_quantities=auxiliary_quantities,
+            truth=truth,
+            graph=graph,
+            theta=theta,
+            phi=phi,
+            pop_synth=pop_synth,
+        )
 
     def to_sub_population(self, observed: bool = True) -> "Population":
         """
         Create a population that is down selected from either the
         observed or unobserved population
 
-        :param observed: extract the observed or unobserved object
-
+        :param observed: Extract the observed or unobserved object
+        :type observed: bool
+        :returns: A new population object
+        :rtype: :class:`Population`
         """
 
         if observed:
@@ -756,8 +772,7 @@ def to_sub_population(self, observed: bool = True) -> "Population":
 
     def display(self):
         """
-        Display the simulation parameters
-
+        Display the simulation parameters.
         """
 
         info = "### %s simulation\nDetected %d out of %d objects" % (
@@ -793,13 +808,11 @@ def display(self):
         display(pd.DataFrame(out))
 
     def display_true_fluxes(self, ax=None, flux_color=dark, **kwargs):
-        """Display the fluxes
-
-        :param ax:
-        :param flux_color:
-        :returns:
-        :rtype:
+        """
+        Display the fluxes.
 
+        :param ax: Axis on which to plot
+        :flux_color: Color of fluxes
         """
 
         if ax is None:
@@ -839,13 +852,10 @@ def display_true_fluxes(self, ax=None, flux_color=dark, **kwargs):
 
     def display_obs_fluxes(self, ax=None, flux_color=dark, **kwargs):
         """
-        display the observed fluxes
-
-        :param ax:
-        :param flux_color:
-        :returns:
-        :rtype:
+        Display the observed fluxes.
 
+        :param ax: Axis on which to plot
+        :flux_color: Color of fluxes
         """
 
         # do not try to plot if there is nothing
@@ -901,16 +911,13 @@ def display_fluxes(
         **kwargs,
     ):
         """
-        display the fluxes
-
-        :param ax:
-        :param true_color:
-        :param obs_color:
-        :param arrow_color:
-        :param with_arrows:
-        :returns:
-        :rtype:
+        Display the fluxes.
 
+        :param ax: Axis on which to plot
+        :param true_color: Color of true values
+        :param obs_color: Color of obs values
+        :param arrow_color: Color of arrows
+        :param with_arrows: If `True`, display arrows
         """
 
         if ax is None:
@@ -958,15 +965,11 @@ def display_luminosities(
         **kwargs,
     ):
         """
-        display the luminosities
-
-        :param ax:
-        :param true_color:
-        :param obs_color:
-
-        :returns:
-        :rtype:
+        Display the luminosities
 
+        :param ax: Axis on which to plot
+        :param true_color: Color of true values
+        :param obs_color: Color of obs values
         """
 
         if ax is None:
@@ -1040,7 +1043,7 @@ def _display_sphere(
             fig.controls = control
             control.autoRotate = True
             fig.render_continuous = True
-            #control.autoRotate = True
+            # control.autoRotate = True
             # toggle_rotate = widgets.ToggleButton(description="Rotate")
             # widgets.jslink((control, "autoRotate"), (toggle_rotate, "value"))
             # r_value = toggle_rotate
@@ -1142,12 +1145,9 @@ def display_flux_sphere(
 
     def display_distances(self, ax=None):
         """
-        display the distances
-
-        :param ax:
-        :returns:
-        :rtype:
+        Display the distances
 
+        :param ax: Axis on which to plot
         """
 
         if ax is None:
diff --git a/popsynth/population_synth.py b/popsynth/population_synth.py
index 6231065..c095b91 100644
--- a/popsynth/population_synth.py
+++ b/popsynth/population_synth.py
@@ -1,4 +1,4 @@
-from abc import ABCMeta, abstractmethod
+from abc import ABCMeta
 from collections import OrderedDict
 from typing import Any, Dict, Optional, Union
 
@@ -12,15 +12,12 @@
 
 # from numpy.typing import np.ndarray
 
-
 from popsynth.auxiliary_sampler import AuxiliarySampler, DerivedLumAuxSampler
 from popsynth.distribution import LuminosityDistribution, SpatialDistribution
 from popsynth.distributions.cosmological_distribution import CosmologicalDistribution
 from popsynth.population import Population
-from popsynth.selection_probability import ( SelectionProbabilty,
-                                             UnitySelection)
+from popsynth.selection_probability import SelectionProbabilty, UnitySelection
 from popsynth.utils.logging import setup_logger
-from popsynth.utils.progress_bar import progress_bar
 from popsynth.utils.registry import (
     auxiliary_parameter_registry,
     distribution_registry,
@@ -42,11 +39,11 @@ def __init__(
         derived luminosity distribution and everything is setup.
 
         :param spatial_distribution: The spatial distribution to sample locations from
-        :param luminosity_distribution: the optional luminosity distribution
-        :param seed: random see
-        :returns:
-        :rtype:
-
+        :type spatial_distribution: :class:`SpatialDistribution`
+        :param luminosity_distribution: The optional luminosity distribution
+        :type luminosity_distribution: :class:`LuminosityDistribution`
+        :param seed: Random seed
+        :type seed: int
         """
 
         self._n_model = 500  # type: int
@@ -117,10 +114,11 @@ def __init__(
 
     def clean(self, reset: bool = False):
         """
-        clean the auxiliary samplers, selections, etc
+        Clean the auxiliary samplers, selections, etc
         from the population synth
-        :param reset: reset any attached distributions and samplers
 
+        :param reset: If `True`, reset any attached distributions and samplers
+        :type reset: bool
         """
 
         if reset:
@@ -176,9 +174,10 @@ def clean(self, reset: bool = False):
 
     def write_to(self, file_name: str) -> None:
         """
-        write the population synth to a file
-        :param file_name: the file name of the output yaml
+        Write the population synth to a YAML file.
 
+        :param file_name: the file name of the output YAML
+        :type file_name: str
         """
         with open(file_name, "w") as f:
 
@@ -191,7 +190,10 @@ def write_to(self, file_name: str) -> None:
 
     def to_dict(self) -> Dict[str, Any]:
         """
-        convert the population synth to a dictionary
+        Convert the population synth to a dictionary
+
+        :returns: Popsynth dict
+        :rtype: Dict[str, Any]
         """
 
         output: Dict[str, Any] = {}
@@ -290,9 +292,12 @@ def to_dict(self) -> Dict[str, Any]:
     @classmethod
     def from_dict(cls, input: Dict[str, Any]) -> "PopulationSynth":
         """
-        build a PopulationSynth object from a dictionary
+        Build a PopulationSynth object from a dictionary
 
         :param input: the dictionary from which to build
+        :type input: Dict[str, Any]
+        :returns: Popsynth object
+        :rtype: :class:`PopulationSynth`
         """
 
         if "luminosity distribution" in input:
@@ -307,7 +312,7 @@ def from_dict(cls, input: Dict[str, Any]) -> "PopulationSynth":
 
             # create the instance
 
-            luminosity_distribtuion: LuminosityDistribution = distribution_registry[
+            luminosity_distribution: LuminosityDistribution = distribution_registry[
                 ld_name]
 
             # now set the values of the parameters
@@ -318,19 +323,19 @@ def from_dict(cls, input: Dict[str, Any]) -> "PopulationSynth":
 
                 log.debug(f"trying to set {k} to {v}")
 
-                for x in luminosity_distribtuion.__class__.mro():
+                for x in luminosity_distribution.__class__.mro():
 
                     if k in x.__dict__:
 
-                        setattr(luminosity_distribtuion, k, float(v))
+                        setattr(luminosity_distribution, k, float(v))
 
                         break
 
-                log.debug(f"{luminosity_distribtuion.params}")
+                log.debug(f"{luminosity_distribution.params}")
 
         else:
 
-            luminosity_distribtuion = None
+            luminosity_distribution = None
 
         # try
 
@@ -363,7 +368,7 @@ def from_dict(cls, input: Dict[str, Any]) -> "PopulationSynth":
 
         pop_synth: PopulationSynth = cls(
             spatial_distribution,
-            luminosity_distribution=luminosity_distribtuion,
+            luminosity_distribution=luminosity_distribution,
             seed=seed,
         )
 
@@ -719,11 +724,10 @@ def add_model_space(self, name, start, stop, log=True):
         """
         Add a model space for stan generated quantities
 
-        :param name: name that Stan will use
-        :param start: start of the grid
-        :param stop: stop of the grid
-        :param log: use log10 or not
-
+        :param name: Name that Stan will use
+        :param start: Start of the grid
+        :param stop: Stop of the grid
+        :param log: Use log10 or not
         """
         if log:
             space = np.logspace(np.log10(start), np.log10(stop), self._n_model)
@@ -738,12 +742,11 @@ def add_auxiliary_sampler(self,
                               auxiliary_sampler: Union[DerivedLumAuxSampler,
                                                        AuxiliarySampler]):
         """
-        add an auxiliary sampler or derived luminosity sampler to the population
-        synth
+        Add an auxiliary sampler or derived luminosity sampler to the population
+        synth.
 
-        :param auxiliary_sampler:
-        :returns:
-        :rtype:
+        :param auxiliary_sampler: The auxiliary_sampler
+        :type auxiliary_sampler: Union[DerivedLumAuxSampler, AuxiliarySampler]
         """
 
         self.add_observed_quantity(auxiliary_sampler)
@@ -752,13 +755,11 @@ def add_observed_quantity(self,
                               auxiliary_sampler: Union[DerivedLumAuxSampler,
                                                        AuxiliarySampler]):
         """
-        add an auxiliary sampler or derived luminosity sampler to the population
+        Add an auxiliary sampler or derived luminosity sampler to the population
         synth
 
-        :param auxiliary_sampler:
-        :returns:
-        :rtype:
-
+        :param auxiliary_sampler: The auxiliary_sampler
+        :type auxiliary_sampler: Union[DerivedLumAuxSampler, AuxiliarySampler]
         """
 
         if isinstance(auxiliary_sampler, DerivedLumAuxSampler):
@@ -801,7 +802,10 @@ def add_observed_quantity(self,
 
     def set_distance_selection(self, selector: SelectionProbabilty) -> None:
         """
-        Set the selection type for the distance
+        Set the selection type for the distance.
+
+        :param selector: The selector
+        :type selector: :class:`SelectionProbabilty`
         """
 
         if not isinstance(selector, SelectionProbabilty):
@@ -816,7 +820,10 @@ def set_distance_selection(self, selector: SelectionProbabilty) -> None:
 
     def set_flux_selection(self, selector: SelectionProbabilty) -> None:
         """
-        Set the selection type for the distance
+        Set the selection type for the flux
+
+        :param selector: The selector
+        :type selector: :class:`SelectionProbabilty`
         """
         if not isinstance(selector, SelectionProbabilty):
 
@@ -832,6 +839,9 @@ def add_spatial_selector(self,
                              spatial_selector: SelectionProbabilty) -> None:
         """
         Add a spatial selector into the mix
+
+        :param spatial_selector: The spatial selector
+        :type spatial_selector: :class:`SelectionProbabilty`
         """
 
         if not isinstance(spatial_selector, SelectionProbabilty):
@@ -848,7 +858,7 @@ def name(self) -> str:
 
     def draw_log10_fobs(self, f, f_sigma, size=1) -> np.ndarray:
         """
-        draw the log10 of the the fluxes
+        Draw the log10 of the the fluxes.
         """
 
         log10_f = np.log10(f)
@@ -861,7 +871,7 @@ def draw_log10_fobs(self, f, f_sigma, size=1) -> np.ndarray:
 
     def draw_log_fobs(self, f, f_sigma, size=1) -> np.ndarray:
         """
-        draw the log10 of the the fluxes
+        Draw the log10 of the the fluxes.
         """
 
         log_f = np.log(f)
@@ -877,17 +887,19 @@ def draw_survey(
         log10_flux_draw: bool = True,
     ) -> Population:
         """
-        Draw the total survey and return a Population object.
-        
+        Draw the total survey and return a :class:`Population` object.
+
         This will sample all attached distributions and apply selection functions.
-        
-        If a value of flux_sigma is given, the log10 observed fluxes are sampled with 
+
+        If a value of flux_sigma is given, the log10 observed fluxes are sampled with
         measurement error.
-        
 
-        :param flux_sigma: the homoskedastic sigma for the flux in log10 space
-        :param log10_flux_draw: if True, fluxes are drawn in log space
+        :param flux_sigma: The homoskedastic sigma for the flux in log10 space
+        :type flux_sigma: Optional[float]
+        :param log10_flux_draw: if `True`, fluxes are drawn in log space
+        :type log10_flux_draw: bool
         :return: a Population object
+        :rtype: :class:`Population`
         """
 
         # this stores all the "true" population values from all the samplers
@@ -1249,8 +1261,7 @@ def draw_survey(
 
     def display(self) -> None:
         """
-        Display the simulation parameters
-
+        Display the simulation parameters.
         """
 
         if self._luminosity_distribution is not None:
@@ -1290,11 +1301,7 @@ def graph(self):
 
     def _build_graph(self):
         """
-        builds the graph for all the samplers
-
-        :returns:
-        :rtype:
-
+        Builds the graph for all the samplers.
         """
 
         # first check out the luminosity sampler
diff --git a/popsynth/populations/bpl_population.py b/popsynth/populations/bpl_population.py
index 267097c..282e1db 100644
--- a/popsynth/populations/bpl_population.py
+++ b/popsynth/populations/bpl_population.py
@@ -8,28 +8,38 @@
 
 
 class BPLHomogeneousSphericalPopulation(SphericalPopulation):
-    def __init__(self,
-                 Lambda,
-                 Lmin,
-                 alpha,
-                 Lbreak,
-                 beta,
-                 Lmax,
-                 r_max=5,
-                 seed=1234):
-        """FIXME! briefly describe function
-
-        :param Lambda:
-        :param Lmin:
-        :param alpha:
-        :param Lbreak:
-        :param beta:
-        :param Lmax:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+    def __init__(
+        self,
+        Lambda: float,
+        Lmin: float,
+        alpha: float,
+        Lbreak: float,
+        beta: float,
+        Lmax: float,
+        r_max: float = 5,
+        seed: int = 1234,
+    ):
+        """
+        A population built on the :class:`ConstantSphericalDistribution`
+        spatial distribution and the :class:`BPLDistribution`
+        luminosity distribution.
+
+        :param Lambda: Density per unit volume
+        :type Lambda: float
+        :param Lmin: Minimum value of the luminosity
+        :type Lmin: float
+        :param alpha: Lower luminosity index
+        :type alpha: float
+        :param Lbreak: Break luminosity
+        :type Lbreak: float
+        :param beta: Upper luminosity index
+        :type beta: float
+        :param Lmax: Maximum value of the luminosity
+        :type Lmax: float
+        :param r_max: Maximum distance
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
         """
         luminosity_distribution = BPLDistribution(seed=seed)
         luminosity_distribution.Lmin = Lmin
@@ -47,30 +57,41 @@ def __init__(self,
 
 
 class BPLZPowerSphericalPopulation(ZPowerSphericalPopulation):
-    def __init__(self,
-                 Lambda,
-                 delta,
-                 Lmin,
-                 alpha,
-                 Lbreak,
-                 beta,
-                 Lmax,
-                 r_max=5,
-                 seed=1234):
-        """FIXME! briefly describe function
-
-        :param Lambda:
-        :param delta:
-        :param Lmin:
-        :param alpha:
-        :param Lbreak:
-        :param beta:
-        :param Lmax:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+    def __init__(
+        self,
+        Lambda: float,
+        delta: float,
+        Lmin: float,
+        alpha: float,
+        Lbreak: float,
+        beta: float,
+        Lmax: float,
+        r_max: float = 5,
+        seed: int = 1234,
+    ):
+        """
+        A population built on the :class:`ZPowerSphericalDistribution`
+        spatial distribution and the :class:`BPLDistribution`
+        luminosity distribution.
+
+        :param Lambda: Density per unit volume
+        :type Lambda: float
+        :param delta: Index of spatial distribution
+        :type delta: float
+        :param Lmin: Minimum value of the luminosity
+        :type Lmin: float
+        :param alpha: Lower luminosity index
+        :type alpha: float
+        :param Lbreak: Break luminosity
+        :type Lbreak: float
+        :param beta: Upper luminosity index
+        :type beta: float
+        :param Lmax: Maximum value of the luminosity
+        :type Lmax: float
+        :param r_max: Maximum distance
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
         """
         luminosity_distribution = BPLDistribution(seed=seed)
         luminosity_distribution.Lmin = Lmin
@@ -91,31 +112,40 @@ def __init__(self,
 class BPLZPowerCosmoPopulation(ZPowerCosmoPopulation):
     def __init__(
         self,
-        Lambda,
-        delta,
-        Lmin,
-        alpha,
-        Lbreak,
-        beta,
-        Lmax,
-        r_max=5,
-        seed=1234,
-        is_rate=True,
+        Lambda: float,
+        delta: float,
+        Lmin: float,
+        alpha: float,
+        Lbreak: float,
+        beta: float,
+        Lmax: float,
+        r_max: float = 5,
+        seed: int = 1234,
+        is_rate: bool = True,
     ):
-        """FIXME! briefly describe function
-
-        :param Lambda:
-        :param delta:
-        :param Lmin:
-        :param alpha:
-        :param Lbreak:
-        :param beta:
-        :param Lmax:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+        """
+        A population built on the :class:`ZPowerCosmoDistribution`
+        spatial distribution and the :class:`BPLDistribution`
+        luminosity distribution.
+
+        :param Lambda: Density in Gpc^-3
+        :type Lambda: float
+        :param delta: Index of spatial distribution
+        :type delta: float
+        :param Lmin: Minimum value of the luminosity
+        :type Lmin: float
+        :param alpha: Lower luminosity index
+        :type alpha: float
+        :param Lbreak: Break luminosity
+        :type Lbreak: float
+        :param beta: Upper luminosity index
+        :type beta: float
+        :param Lmax: Maximum value of the luminosity
+        :type Lmax: float
+        :param r_max: Maximum distance
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
         """
 
         luminosity_distribution = BPLDistribution(seed=seed)
@@ -138,37 +168,49 @@ def __init__(
 class BPLSFRPopulation(SFRPopulation):
     def __init__(
         self,
-        r0,
-        a,
-        rise,
-        decay,
-        peak,
-        Lmin,
-        alpha,
-        Lbreak,
-        beta,
-        Lmax,
-        r_max=5,
-        seed=1234,
-        is_rate=True,
+        r0: float,
+        a: float,
+        rise: float,
+        decay: float,
+        peak: float,
+        Lmin: float,
+        alpha: float,
+        Lbreak: float,
+        beta: float,
+        Lmax: float,
+        r_max: float = 5,
+        seed: int = 1234,
+        is_rate: bool = True,
     ):
-        """FIXME! briefly describe function
-
-        :param r0:
-        :param a:
-        :param rise:
-        :param decay:
-        :param peak:
-        :param Lmin:
-        :param alpha:
-        :param Lbreak:
-        :param beta:
-        :param Lmax:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+        """
+        A population built on the :class:`SFRDistribution`
+        spatial distribution and the :class:`BPLDistribution`
+        luminosity distribution.
+
+        :param r0: Local density in units of Gpc^-3
+        :type r0: float
+        :param a: Offset at z=0
+        :type a: float
+        :param rise: Rise at low z
+        :type rise: float
+        :param decay: Decay at high z
+        :type decay: float
+        :param peak: Peak of z distribution
+        :type peak: float
+        :param Lmin: Minimum value of the luminosity
+        :type Lmin: float
+        :param alpha: Lower luminosity index
+        :type alpha: float
+        :param Lbreak: Break luminosity
+        :type Lbreak: float
+        :param beta: Upper luminosity index
+        :type beta: float
+        :param Lmax: Maximum value of the luminosity
+        :type Lmax: float
+        :param r_max: Maximum redshift
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
         """
 
         luminosity_distribution = BPLDistribution(seed=seed)
diff --git a/popsynth/populations/lognormal_population.py b/popsynth/populations/lognormal_population.py
index 48d53b6..1e38b57 100644
--- a/popsynth/populations/lognormal_population.py
+++ b/popsynth/populations/lognormal_population.py
@@ -10,17 +10,29 @@
 
 
 class LogNormalHomogeneousSphericalPopulation(SphericalPopulation):
-    def __init__(self, Lambda, mu, tau, r_max=5, seed=1234):
-        """FIXME! briefly describe function
-
-        :param Lambda:
-        :param mu:
-        :param tau:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+    def __init__(
+        self,
+        Lambda: float,
+        mu: float,
+        tau: float,
+        r_max: float = 5,
+        seed: int = 1234,
+    ):
+        """
+        A population built on the :class:`ConstantSphericalDistribution`
+        spatial distribution and the :class:`LogNormalDistribution`
+        luminosity distribution.
+
+        :param Lambda: Density per unit volume
+        :type Lambda: float
+        :param mu: Mean of the luminosity distribution
+        :type mu: float
+        :param tau: Standard deviation of the luminosity distribution
+        :type tau: float
+        :param r_max: Maximum distance
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
         """
 
         luminosity_distribution = LogNormalDistribution(seed=seed)
@@ -36,18 +48,32 @@ def __init__(self, Lambda, mu, tau, r_max=5, seed=1234):
 
 
 class LogNormalZPowerSphericalPopulation(ZPowerSphericalPopulation):
-    def __init__(self, Lambda, delta, mu, tau, r_max=5, seed=1234):
-        """FIXME! briefly describe function
-
-        :param Lambda:
-        :param delta:
-        :param mu:
-        :param tau:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+    def __init__(
+        self,
+        Lambda: float,
+        delta: float,
+        mu: float,
+        tau: float,
+        r_max: float = 5,
+        seed: int = 1234,
+    ):
+        """
+        A population built on the :class:`ZPowerSphericalDistribution`
+        spatial distribution and the :class:`LogNormalDistribution`
+        luminosity distribution.
+
+        :param Lambda: Density per unit volume
+        :type Lambda: float
+        :param delta: Index of the spatial distribution
+        :type delta: float
+        :param mu: Mean of the luminosity distribution
+        :type mu: float
+        :param tau: Standard deviation of the luminosity distribution
+        :type tau: float
+        :param r_max: Maximum distance
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
         """
 
         luminosity_distribution = LogNormalDistribution(seed=seed)
@@ -66,25 +92,36 @@ def __init__(self, Lambda, delta, mu, tau, r_max=5, seed=1234):
 class LogNormalZPowerCosmoPopulation(ZPowerCosmoPopulation):
     def __init__(
         self,
-        Lambda,
-        delta,
-        mu,
-        tau,
-        r_max=5,
-        seed=1234,
-        is_rate=True,
+        Lambda: float,
+        delta: float,
+        mu: float,
+        tau: float,
+        r_max: float = 5,
+        seed: int = 1234,
+        is_rate: bool = True,
     ):
-        """FIXME! briefly describe function
-
-        :param Lambda:
-        :param delta:
-        :param mu:
-        :param tau:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+        """
+        A population built on the :class:`ZPowerCosmoDistribution`
+        spatial distribution and the :class:`LogNormalDistribution`
+        luminosity distribution.
+
+        :param Lambda: Density in units of Gpc^-3
+        :type Lambda: float
+        :param delta: Index of the spatial distribution
+        :type delta: float
+        :param mu: Mean of the luminosity distribution
+        :type mu: float
+        :param tau: Standard deviation of the luminosity distribution
+        :type tau: float
+        :param r_max: Maximum redshift
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
+        :param is_rate: `True` if modelling a population of transient events,
+        `False` if modelling a population of steady-state objects.
+        Affects the ``time_adjustment`` method used in cosmo calculations.
+        Default is `True`.
+        :type is_rate: bool
         """
 
         luminosity_distribution = LogNormalDistribution(seed=seed)
@@ -102,31 +139,47 @@ def __init__(
 
 
 class LogNormalSFRPopulation(SFRPopulation):
-    def __init__(self,
-                 r0,
-                 a,
-                 rise,
-                 decay,
-                 peak,
-                 mu,
-                 tau,
-                 r_max=5,
-                 seed=1234,
-                 is_rate=True):
-        """FIXME! briefly describe function
-
-        :param r0:
-        :param a:
-        :param rise:
-        :param decay:
-        :param peak:
-         :param mu:
-        :param tau:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+    def __init__(
+        self,
+        r0: float,
+        a: float,
+        rise: float,
+        decay: float,
+        peak: float,
+        mu: float,
+        tau: float,
+        r_max: float = 5,
+        seed: int = 1234,
+        is_rate: bool = True,
+    ):
+        """
+        A population built on the :class:`SFRDistribution`
+        spatial distribution and the :class:`LogNormalDistribution`
+        luminosity distribution.
+
+        :param r0: Local density in units of Gpc^-3
+        :type r0: float
+        :param a: Offset at z=0
+        :type a: float
+        :param rise: Rise at low z
+        :type rise: float
+        :param decay: Decay at high z
+        :type decay: float
+        :param peak: Peak of z distribution
+        :type peak: float
+        :param mu: Mean of the luminosity distribution
+        :type mu: float
+        :param tau: Standard deviation of the luminosity distribution
+        :type tau: float
+        :param r_max: Maximum redshift
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
+        :param is_rate: `True` if modelling a population of transient events,
+        `False` if modelling a population of steady-state objects.
+        Affects the ``time_adjustment`` method used in cosmo calculations.
+        Default is `True`.
+        :type is_rate: bool
         """
 
         luminosity_distribution = LogNormalDistribution(seed=seed)
@@ -147,17 +200,29 @@ def __init__(self,
 
 
 class Log10NormalHomogeneousSphericalPopulation(SphericalPopulation):
-    def __init__(self, Lambda, mu, tau, r_max=5, seed=1234):
-        """FIXME! briefly describe function
-
-        :param Lambda:
-        :param mu:
-        :param tau:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+    def __init__(
+        self,
+        Lambda: float,
+        mu: float,
+        tau: float,
+        r_max: float = 5,
+        seed: int = 1234,
+    ):
+        """
+        A population built on the :class:`ConstantSphericalDistribution`
+        spatial distribution and the :class:`Log10NormalDistribution`
+        luminosity distribution.
+
+        :param Lambda: Density per unit volume
+        :type Lambda: float
+        :param mu: Mean of the luminosity distribution
+        :type mu: float
+        :param tau: Standard deviation of the luminosity distribution
+        :type tau: float
+        :param r_max: Maximum distance
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
         """
 
         luminosity_distribution = Log10NormalDistribution(seed=seed)
@@ -173,18 +238,32 @@ def __init__(self, Lambda, mu, tau, r_max=5, seed=1234):
 
 
 class Log10NormalZPowerSphericalPopulation(ZPowerSphericalPopulation):
-    def __init__(self, Lambda, delta, mu, tau, r_max=5, seed=1234):
-        """FIXME! briefly describe function
-
-        :param Lambda:
-        :param delta:
-        :param mu:
-        :param tau:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+    def __init__(
+        self,
+        Lambda: float,
+        delta: float,
+        mu: float,
+        tau: float,
+        r_max: float = 5,
+        seed: int = 1234,
+    ):
+        """
+        A population built on the :class:`ZPowerSphericalDistribution`
+        spatial distribution and the :class:`Log10NormalDistribution`
+        luminosity distribution.
+
+        :param Lambda: Density per unit volume
+        :type Lambda: float
+        :param delta: Index of the spatial distribution
+        :type delta: float
+        :param mu: Mean of the luminosity distribution
+        :type mu: float
+        :param tau: Standard deviation of the luminosity distribution
+        :type tau: float
+        :param r_max: Maximum distance
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
         """
 
         luminosity_distribution = Log10NormalDistribution(seed=seed)
@@ -203,25 +282,36 @@ def __init__(self, Lambda, delta, mu, tau, r_max=5, seed=1234):
 class Log10NormalZPowerCosmoPopulation(ZPowerCosmoPopulation):
     def __init__(
         self,
-        Lambda,
-        delta,
-        mu,
-        tau,
-        r_max=5,
-        seed=1234,
-        is_rate=True,
+        Lambda: float,
+        delta: float,
+        mu: float,
+        tau: float,
+        r_max: float = 5,
+        seed: float = 1234,
+        is_rate: float = True,
     ):
-        """FIXME! briefly describe function
-
-        :param Lambda:
-        :param delta:
-        :param mu:
-        :param tau:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+        """
+        A population built on the :class:`ZPowerCosmoDistribution`
+        spatial distribution and the :class:`Log10NormalDistribution`
+        luminosity distribution.
+
+        :param Lambda: Density in units of Gpc^-3
+        :type Lambda: float
+        :param delta: Index of the spatial distribution
+        :type delta: float
+        :param mu: Mean of the luminosity distribution
+        :type mu: float
+        :param tau: Standard deviation of the luminosity distribution
+        :type tau: float
+        :param r_max: Maximum redshift
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
+        :param is_rate: `True` if modelling a population of transient events,
+        `False` if modelling a population of steady-state objects.
+        Affects the ``time_adjustment`` method used in cosmo calculations.
+        Default is `True`.
+        :type is_rate: bool
         """
 
         luminosity_distribution = Log10NormalDistribution(seed=seed)
@@ -241,31 +331,45 @@ def __init__(
 class Log10NormalSFRPopulation(SFRPopulation):
     def __init__(
         self,
-        r0,
-        a,
-        rise,
-        decay,
-        peak,
-        mu,
-        tau,
-        r_max=5,
-        seed=1234,
-        is_rate=True,
+        r0: float,
+        a: float,
+        rise: float,
+        decay: float,
+        peak: float,
+        mu: float,
+        tau: float,
+        r_max: float = 5,
+        seed: int = 1234,
+        is_rate: bool = True,
     ):
-        """FIXME! briefly describe function
-
-        :param r0:
-        :param a:
-        :param rise:
-        :param decay:
-        :param peak:
-        :param mu:
-        :param tau:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+        """
+        A population built on the :class:`SFRDistribution`
+        spatial distribution and the :class:`Log10NormalDistribution`
+        luminosity distribution.
+
+        :param r0: Local density in units of Gpc^-3
+        :type r0: float
+        :param a: Offset at z=0
+        :type a: float
+        :param rise: Rise at low z
+        :type rise: float
+        :param decay: Decay at high z
+        :type decay: float
+        :param peak: Peak of z distribution
+        :type peak: float
+        :param mu: Mean of the luminosity distribution
+        :type mu: float
+        :param tau: Standard deviation of the luminosity distribution
+        :type tau: float
+        :param r_max: Maximum redshift
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
+        :param is_rate: `True` if modelling a population of transient events,
+        `False` if modelling a population of steady-state objects.
+        Affects the ``time_adjustment`` method used in cosmo calculations.
+        Default is `True`.
+        :type is_rate: bool
         """
 
         luminosity_distribution = Log10NormalDistribution(seed=seed)
diff --git a/popsynth/populations/pareto_populations.py b/popsynth/populations/pareto_populations.py
index a65cdb2..5e1b5d3 100644
--- a/popsynth/populations/pareto_populations.py
+++ b/popsynth/populations/pareto_populations.py
@@ -8,17 +8,29 @@
 
 
 class ParetoHomogeneousSphericalPopulation(SphericalPopulation):
-    def __init__(self, Lambda, Lmin, alpha, r_max=5, seed=1234):
-        """FIXME! briefly describe function
-
-        :param Lambda:
-        :param Lmin:
-        :param alpha:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+    def __init__(
+        self,
+        Lambda: float,
+        Lmin: float,
+        alpha: float,
+        r_max: float = 5,
+        seed: int = 1234,
+    ):
+        """
+        A population built on the :class:`ConstantSphericalDistribution`
+        spatial distribution and the :class:`ParetoDistribution`
+        luminosity distribution.
+
+        :param Lambda: Density per unit volume
+        :type Lambda: float
+        :param Lmin: Minimum value of the luminosity
+        :type Lmin: float
+        :param alpha: Index of the luminosity distribution
+        :type alpha: float
+        :param r_max: Maximum distance
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
         """
 
         luminosity_distribution = ParetoDistribution(seed=seed)
@@ -34,18 +46,32 @@ def __init__(self, Lambda, Lmin, alpha, r_max=5, seed=1234):
 
 
 class ParetoZPowerSphericalPopulation(ZPowerSphericalPopulation):
-    def __init__(self, Lambda, delta, Lmin, alpha, r_max=5, seed=1234):
-        """FIXME! briefly describe function
-
-        :param Lambda:
-        :param delta:
-        :param Lmin:
-        :param alpha:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+    def __init__(
+        self,
+        Lambda: float,
+        delta: float,
+        Lmin: float,
+        alpha: float,
+        r_max: float = 5,
+        seed: int = 1234,
+    ):
+        """
+        A population built on the :class:`ZPowerSphericalDistribution`
+        spatial distribution and the :class:`ParetoDistribution`
+        luminosity distribution.
+
+        :param Lambda: Density per unit volume
+        :type Lambda: float
+        :param delta: Index of the spatial distribution
+        :type delta: float
+        :param Lmin: Minimum value of the luminosity
+        :type Lmin: float
+        :param alpha: Index of the luminosity distribution
+        :type alpha: float
+        :param r_max: Maximum distance
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
         """
 
         luminosity_distribution = ParetoDistribution(seed=seed)
@@ -64,25 +90,36 @@ def __init__(self, Lambda, delta, Lmin, alpha, r_max=5, seed=1234):
 class ParetoZPowerCosmoPopulation(ZPowerCosmoPopulation):
     def __init__(
         self,
-        Lambda,
-        delta,
-        Lmin,
-        alpha,
-        r_max=5,
-        seed=1234,
-        is_rate=True,
+        Lambda: float,
+        delta: float,
+        Lmin: float,
+        alpha: float,
+        r_max: float = 5,
+        seed: int = 1234,
+        is_rate: bool = True,
     ):
-        """FIXME! briefly describe function
-
-        :param Lambda:
-        :param delta:
-        :param Lmin:
-        :param alpha:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+        """
+        A population built on the :class:`ZPowerCosmoDistribution`
+        spatial distribution and the :class:`ParetoDistribution`
+        luminosity distribution.
+
+        :param Lambda: Density in units of Gpc^-3
+        :type Lambda: float
+        :param delta: Index of the spatial distribution
+        :type delta: float
+        :param Lmin: Minimum value of the luminosity
+        :type Lmin: float
+        :param alpha: Index of the luminosity distribution
+        :type alpha: float
+        :param r_max: Maximum redshift
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
+        :param is_rate: `True` if modelling a population of transient events,
+        `False` if modelling a population of steady-state objects.
+        Affects the ``time_adjustment`` method used in cosmo calculations.
+        Default is `True`.
+        :type is_rate: bool
         """
 
         luminosity_distribution = ParetoDistribution(seed=seed)
@@ -102,31 +139,45 @@ def __init__(
 class ParetoSFRPopulation(SFRPopulation):
     def __init__(
         self,
-        r0,
-        a,
-        rise,
-        decay,
-        peak,
-        Lmin,
-        alpha,
-        r_max=5,
-        seed=1234,
-        is_rate=True,
+        r0: float,
+        a: float,
+        rise: float,
+        decay: float,
+        peak: float,
+        Lmin: float,
+        alpha: float,
+        r_max: float = 5,
+        seed: int = 1234,
+        is_rate: bool = True,
     ):
-        """FIXME! briefly describe function
-
-        :param r0:
-        :param a:
-        :param rise:
-        :param decay:
-        :param peak:
-        :param Lmin:
-        :param alpha:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+        """
+        A population built on the :class:`SFRDistribution`
+        spatial distribution and the :class:`ParetoDistribution`
+        luminosity distribution.
+
+        :param r0: Local density in units of Gpc^-3
+        :type r0: float
+        :param a: Offset at z=0
+        :type a: float
+        :param rise: Rise at low z
+        :type rise: float
+        :param decay: Decay at high z
+        :type decay: float
+        :param peak: Peak of z distribution
+        :type peak: float
+        :param Lmin: Minimum value of the luminosity
+        :type Lmin: float
+        :param alpha: Index of the luminosity distribution
+        :type alpha: float
+        :param r_max: Maximum redshift
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
+        :param is_rate: `True` if modelling a population of transient events,
+        `False` if modelling a population of steady-state objects.
+        Affects the ``time_adjustment`` method used in cosmo calculations.
+        Default is `True`.
+        :type is_rate: bool
         """
 
         luminosity_distribution = ParetoDistribution(seed=seed)
diff --git a/popsynth/populations/schechter_populations.py b/popsynth/populations/schechter_populations.py
index 034e448..99f2dae 100644
--- a/popsynth/populations/schechter_populations.py
+++ b/popsynth/populations/schechter_populations.py
@@ -9,17 +9,29 @@
 
 
 class SchechterHomogeneousSphericalPopulation(SphericalPopulation):
-    def __init__(self, Lambda, Lmin, alpha, r_max=10, seed=1234):
-        """FIXME! briefly describe function
-
-        :param Lambda:
-        :param Lmin:
-        :param alpha:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+    def __init__(
+        self,
+        Lambda: float,
+        Lmin: float,
+        alpha: float,
+        r_max: float = 10,
+        seed: int = 1234,
+    ):
+        """
+        A population built on the :class:`ConstantSphericalDistribution`
+        spatial distribution and the :class:`SchechterDistribution`
+        luminosity distribution.
+
+        :param Lambda: Density per unit volume
+        :type Lambda: float
+        :param Lmin: Minimum value of the luminosity
+        :type Lmin: float
+        :param alpha: Index of the luminosity distribution
+        :type alpha: float
+        :param r_max: Maximum distance
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
         """
 
         luminosity_distribution = SchechterDistribution(seed=seed)
@@ -35,18 +47,32 @@ def __init__(self, Lambda, Lmin, alpha, r_max=10, seed=1234):
 
 
 class SchechterZPowerSphericalPopulation(ZPowerSphericalPopulation):
-    def __init__(self, Lambda, delta, Lmin, alpha, r_max=10, seed=1234):
-        """FIXME! briefly describe function
-
-        :param Lambda:
-        :param delta:
-        :param Lmin:
-        :param alpha:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+    def __init__(
+        self,
+        Lambda: float,
+        delta: float,
+        Lmin: float,
+        alpha: float,
+        r_max: float = 10,
+        seed: int = 1234,
+    ):
+        """
+        A population built on the :class:`ZPowerSphericalDistribution`
+        spatial distribution and the :class:`SchechterDistribution`
+        luminosity distribution.
+
+        :param Lambda: Density per unit volume
+        :type Lambda: float
+        :param delta: Index of the spatial distribution
+        :type delta: float
+        :param Lmin: Minimum value of the luminosity
+        :type Lmin: float
+        :param alpha: Index of the luminosity distribution
+        :type alpha: float
+        :param r_max: Maximum distance
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
         """
 
         luminosity_distribution = SchechterDistribution(seed=seed)
@@ -66,25 +92,36 @@ def __init__(self, Lambda, delta, Lmin, alpha, r_max=10, seed=1234):
 class SchechterZPowerCosmoPopulation(ZPowerCosmoPopulation):
     def __init__(
         self,
-        Lambda,
-        delta,
-        Lmin,
-        alpha,
-        r_max=10,
-        seed=1234,
-        is_rate=True,
+        Lambda: float,
+        delta: float,
+        Lmin: float,
+        alpha: float,
+        r_max: float = 10,
+        seed: int = 1234,
+        is_rate: bool = True,
     ):
-        """FIXME! briefly describe function
-
-        :param Lambda:
-        :param delta:
-        :param Lmin:
-        :param alpha:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+        """
+        A population built on the :class:`ZPowerCosmoDistribution`
+        spatial distribution and the :class:`SchechterDistribution`
+        luminosity distribution.
+
+        :param Lambda: Density in units of Gpc^-3
+        :type Lambda: float
+        :param delta: Index of the spatial distribution
+        :type delta: float
+        :param Lmin: Minimum value of the luminosity
+        :type Lmin: float
+        :param alpha: Index of the luminosity distribution
+        :type alpha: float
+        :param r_max: Maximum redshift
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
+        :param is_rate: `True` if modelling a population of transient events,
+        `False` if modelling a population of steady-state objects.
+        Affects the ``time_adjustment`` method used in cosmo calculations.
+        Default is `True`.
+        :type is_rate: bool
         """
 
         luminosity_distribution = SchechterDistribution(seed=seed)
@@ -105,31 +142,45 @@ def __init__(
 class SchechterSFRPopulation(SFRPopulation):
     def __init__(
         self,
-        r0,
-        a,
-        rise,
-        decay,
-        peak,
-        Lmin,
-        alpha,
-        r_max=10,
-        seed=1234,
-        is_rate=True,
+        r0: float,
+        a: float,
+        rise: float,
+        decay: float,
+        peak: float,
+        Lmin: float,
+        alpha: float,
+        r_max: float = 10,
+        seed: int = 1234,
+        is_rate: bool = True,
     ):
-        """FIXME! briefly describe function
-
-        :param r0:
-        :param a:
-        :param rise:
-        :param decay:
-        :param peak:
-        :param Lmin:
-        :param alpha:
-        :param r_max:
-        :param seed:
-        :returns:
-        :rtype:
-
+        """
+        A population built on the :class:`SFRDistribution`
+        spatial distribution and the :class:`SchechterDistribution`
+        luminosity distribution.
+
+        :param r0: Local density in units of Gpc^-3
+        :type r0: float
+        :param a: Offset at z=0
+        :type a: float
+        :param rise: Rise at low z
+        :type rise: float
+        :param decay: Decay at high z
+        :type decay: float
+        :param peak: Peak of z distribution
+        :type peak: float
+        :param Lmin: Minimum value of the luminosity
+        :type Lmin: float
+        :param alpha: Index of the luminosity distribution
+        :type alpha: float
+        :param r_max: Maximum redshift
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
+        :param is_rate: `True` if modelling a population of transient events,
+        `False` if modelling a population of steady-state objects.
+        Affects the ``time_adjustment`` method used in cosmo calculations.
+        Default is `True`.
+        :type is_rate: bool
         """
 
         luminosity_distribution = SchechterDistribution(seed=seed)
diff --git a/popsynth/populations/spatial_populations.py b/popsynth/populations/spatial_populations.py
index f4e7ec0..3d0237d 100644
--- a/popsynth/populations/spatial_populations.py
+++ b/popsynth/populations/spatial_populations.py
@@ -1,30 +1,39 @@
 from popsynth.distributions.cosmological_distribution import (  # MergerDistribution,
-    SFRDistribution, ZPowerCosmoDistribution)
+    SFRDistribution, ZPowerCosmoDistribution,
+)
 from popsynth.distributions.spherical_distribution import (
-    ConstantSphericalDistribution, ZPowerSphericalDistribution)
-from popsynth.distributions.spiral_galaxy_distribution import \
-    SpiralGalaxyDistribution
+    ConstantSphericalDistribution,
+    ZPowerSphericalDistribution,
+)
+from popsynth.distribution import LuminosityDistribution
+from popsynth.distributions.spiral_galaxy_distribution import SpiralGalaxyDistribution
 from popsynth.population_synth import PopulationSynth
-
-# First create a bunch of spatial distributions
+"""
+Create a range of spatial populations that can
+be expanded with luminosity distributions.
+"""
 
 
 class SphericalPopulation(PopulationSynth):
-    def __init__(self,
-                 Lambda,
-                 r_max=5.0,
-                 seed=1234,
-                 luminosity_distribution=None):
+    def __init__(
+        self,
+        Lambda: float,
+        r_max: float = 5.0,
+        seed: int = 1234,
+        luminosity_distribution: LuminosityDistribution = None,
+    ):
         """
-        A generic spherical population
-
-        :param Lambda:
-        :param r_max:
-        :param seed:
-        :param luminosity_distribution:
-        :returns:
-        :rtype:
-
+        A generic spherical population based on
+        :class:`ConstantSphericalDistribution`.
+
+        :param Lambda: Density per unit volume
+        :type Lambda: float
+        :param r_max: Maximum distance
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
+        :param luminosity_distribution: Luminosity distribution
+        :type luminosity_distribution: :class:`LuminosityDistribution`, optional
         """
 
         spatial_distribution = ConstantSphericalDistribution(seed=seed)
@@ -39,25 +48,29 @@ def __init__(self,
 
 
 class ZPowerSphericalPopulation(PopulationSynth):
-    def __init__(self,
-                 Lambda,
-                 delta,
-                 r_max=5.0,
-                 seed=1234,
-                 luminosity_distribution=None):
+    def __init__(
+        self,
+        Lambda: float,
+        delta: float,
+        r_max: float = 5.0,
+        seed: int = 1234,
+        luminosity_distribution: LuminosityDistribution = None,
+    ):
         """
-        
-        A spherical population with a density that 
-        scales as (r+1)^delta
-
-        :param Lambda:
-        :param delta:
-        :param r_max:
-        :param seed:
-        :param luminosity_distribution:
-        :returns:
-        :rtype:
-
+        A spherical population with a density that
+        scales as (r+1)^delta.
+        Based on :class:`ZpowerSphericalDistribution`.
+
+        :param Lambda: Local density per unit volume
+        :type Lambda: float
+        :param delta: Index of spatial distribution
+        :type delta: float
+        :param r_max: Maximum distance
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
+        :param luminosity_distribution: Luminosity distribution
+        :type luminosity_distribution: :class:`LuminosityDistribution`, optional
         """
 
         spatial_distribution = ZPowerSphericalDistribution(seed=seed)
@@ -76,27 +89,39 @@ def __init__(self,
 class ZPowerCosmoPopulation(PopulationSynth):
     def __init__(
         self,
-        Lambda,
-        delta,
-        r_max=5.0,
-        seed=1234,
-        luminosity_distribution=None,
-        is_rate=True,
+        Lambda: float,
+        delta: float,
+        r_max: float = 5.0,
+        seed: int = 1234,
+        luminosity_distribution: LuminosityDistribution = None,
+        is_rate: bool = True,
     ):
-        """FIXME! briefly describe function
-
-        :param Lambda:
-        :param delta:
-        :param r_max:
-        :param seed:
-        :param luminosity_distribution:
-        :returns:
-        :rtype:
-
+        """
+        A cosmological population with a density that
+        scales as (z+1)^delta.
+        Based on :class:`ZpowerCosmoDistribution`.
+
+        :param Lambda: Local density in units of Gpc^-3
+        :type Lambda: float
+        :param delta: Index of spatial distribution
+        :type delta: float
+        :param r_max: Maximum redshift
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
+        :param luminosity_distribution: Luminosity distribution
+        :type luminosity_distribution: :class:`LuminosityDistribution`, optional
+        :param is_rate: `True` if modelling a population of transient events,
+        `False` if modelling a population of steady-state objects.
+        Affects the ``time_adjustment`` method used in cosmo calculations.
+        Default is `True`.
+        :type is_rate: bool
         """
 
-        spatial_distribution = ZPowerCosmoDistribution(seed=seed,
-                                                       is_rate=is_rate)
+        spatial_distribution = ZPowerCosmoDistribution(
+            seed=seed,
+            is_rate=is_rate,
+        )
 
         spatial_distribution.Lambda = Lambda
         spatial_distribution.r_max = r_max
@@ -119,22 +144,35 @@ def __init__(
         peak: float,
         r_max: float = 5,
         seed: int = 1234,
-        luminosity_distribution=None,
-        is_rate=True,
+        luminosity_distribution: LuminosityDistribution = None,
+        is_rate: bool = True,
     ):
-        """FIXME! briefly describe function
-
-        :param r0:
-        :param a:
-        :param rise:
-        :param decay:
-        :param peak:
-        :param r_max:
-        :param seed:
-        :param luminosity_distribution:
-        :returns:
-        :rtype:
-
+        """
+        A cosmological population with a density that
+        scales similarly to the star formation rate.
+        Based on :class:`ZpowerCosmoDistribution`.
+
+        :param r0: Local density in units of Gpc^-3
+        :type r0: float
+        :param a: Offset at z=0
+        :type a: float
+        :param rise: Rise at low z
+        :type rise: float
+        :param decay: Decay at high z
+        :type decay: float
+        :param peak: Peak of z distribution
+        :type peak: float
+        :param r_max: Maximum redshift
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
+        :param luminosity_distribution: Luminosity distribution
+        :type luminosity_distribution: :class:`LuminosityDistribution`, optional
+        :param is_rate: `True` if modelling a population of transient events,
+        `False` if modelling a population of steady-state objects.
+        Affects the ``time_adjustment`` method used in cosmo calculations.
+        Default is `True`.
+        :type is_rate: bool
         """
 
         spatial_distribution = SFRDistribution(seed=seed, is_rate=is_rate)
@@ -156,15 +194,41 @@ def __init__(
 class MWRadialPopulation(PopulationSynth):
     def __init__(
         self,
-        rho,
-        a=1.64,
-        b=4.01,
-        R1=0.55,
-        R0=8.5,
-        r_max=20,
-        seed=1234,
-        luminosity_distribution=None,
+        rho: float,
+        a: float = 1.64,
+        b: float = 4.01,
+        R1: float = 0.55,
+        R0: float = 8.5,
+        r_max: float = 20,
+        seed: int = 1234,
+        luminosity_distribution: LuminosityDistribution = None,
     ):
+        """
+        A Milky-way like population based on
+        :class:`SpiralGalaxyDistribution`.
+
+        :param rho: Local density
+        :type rho: float
+        :param a: Shape parameter
+        :type a: float
+        :param b: Shape parameter
+        :type b: float
+        :param R1: Scale parameter
+        :type R1: float
+        :param R0: Scale parameter
+        :type R0: float
+        :param r_max: Maximum distance
+        :type r_max: float
+        :param seed: Random seed
+        :type seed: int
+        :param luminosity_distribution: Luminosity distribution
+        :type luminosity_distribution: :class:`LuminosityDistribution`, optional
+        :param is_rate: `True` if modelling a population of transient events,
+        `False` if modelling a population of steady-state objects.
+        Affects the ``time_adjustment`` method used in cosmo calculations.
+        Default is `True`.
+        :type is_rate: bool
+        """
 
         spatial_distribution = SpiralGalaxyDistribution(seed=seed)
 
diff --git a/popsynth/selection_probability/__init__.py b/popsynth/selection_probability/__init__.py
index 0dd9eef..d335d0c 100644
--- a/popsynth/selection_probability/__init__.py
+++ b/popsynth/selection_probability/__init__.py
@@ -1,13 +1,31 @@
 from .flux_selectors import HardFluxSelection, SoftFluxSelection
-from .generic_selectors import (BernoulliSelection, BoxSelection, LowerBound,
-                                SoftSelection, UnitySelection, UpperBound)
-from .selection_probability import (DummySelection, SelectionParameter,
-                                    SelectionProbabilty)
-from .spatial_selection import SpatialSelection
+from .generic_selectors import (
+    BernoulliSelection,
+    BoxSelection,
+    LowerBound,
+    SoftSelection,
+    UnitySelection,
+    UpperBound,
+)
+from .selection_probability import (
+    DummySelection,
+    SelectionParameter,
+    SelectionProbabilty,
+)
+from .spatial_selection import SpatialSelection, GalacticPlaneSelection
 
 __all__ = [
-    "HardFluxSelection", "SoftFluxSelection", "BernoulliSelection",
-    "LowerBound", "UpperBound", "SoftSelection", "UnitySelection",
-    "SelectionProbabilty", "SpatialSelection", "BoxSelection",
-    "DummySelection", "SelectionParameter"
+    "HardFluxSelection",
+    "SoftFluxSelection",
+    "BernoulliSelection",
+    "LowerBound",
+    "UpperBound",
+    "SoftSelection",
+    "UnitySelection",
+    "SelectionProbabilty",
+    "SpatialSelection",
+    "BoxSelection",
+    "DummySelection",
+    "SelectionParameter",
+    "GalacticPlaneSelection",
 ]
diff --git a/popsynth/selection_probability/flux_selectors.py b/popsynth/selection_probability/flux_selectors.py
index b1aa29c..c71dc24 100644
--- a/popsynth/selection_probability/flux_selectors.py
+++ b/popsynth/selection_probability/flux_selectors.py
@@ -1,5 +1,3 @@
-import numpy as np
-
 from popsynth.utils.logging import setup_logger
 
 from .generic_selectors import LowerBound, SoftSelection
@@ -11,6 +9,11 @@ class HardFluxSelection(LowerBound):
     _selection_name = "HardFluxSelection"
 
     def __init__(self) -> None:
+        """
+        A hard selection on the observed flux.
+
+        Based on :class:`LowerBound`.
+        """
 
         super(HardFluxSelection, self).__init__(use_flux=True)
 
@@ -23,6 +26,11 @@ class SoftFluxSelection(SoftSelection):
     _selection_name = "SoftFluxSelection"
 
     def __init__(self) -> None:
+        """
+        A soft selection on the observed flux.
+
+        Based on :class:`SoftSelection`.
+        """
 
         super(SoftFluxSelection, self).__init__(use_flux=True)
 
diff --git a/popsynth/selection_probability/generic_selectors.py b/popsynth/selection_probability/generic_selectors.py
index 5fe4669..541c4c3 100644
--- a/popsynth/selection_probability/generic_selectors.py
+++ b/popsynth/selection_probability/generic_selectors.py
@@ -16,7 +16,9 @@ class UnitySelection(SelectionProbabilty):
 
     def __init__(self, name="unity"):
         """
-        A selection that returns all unity
+        A selection that returns all selected.
+
+        :param name: Name of the selection
         """
         super(UnitySelection, self).__init__(name=name)
 
@@ -31,6 +33,12 @@ class BernoulliSelection(SelectionProbabilty):
     probability = SelectionParameter(vmin=0, vmax=1, default=0.5)
 
     def __init__(self) -> None:
+        """
+        A Bernoulli selection with ``probability`` as a parameter.
+
+        :param probability: Probability for each Bernoulli trial
+        :type probability: :class:`SelectionParameter`
+        """
 
         super(BernoulliSelection, self).__init__(name="Bernoulli")
 
@@ -60,18 +68,44 @@ class BoxSelection(SelectionProbabilty):
     vmin = SelectionParameter()
     vmax = SelectionParameter()
 
-    def __init__(self,
-                 name: str = "box selection",
-                 use_obs_value: bool = False,
-                 use_distance=False,
-                 use_luminosity=False,
-                 use_flux: bool = False):
+    def __init__(
+        self,
+        name: str = "box selection",
+        use_obs_value: bool = False,
+        use_distance: bool = False,
+        use_luminosity: bool = False,
+        use_flux: bool = False,
+    ):
+        """
+        A box selection on observed_value, distance,
+        luminosity or flux.
+
+        :param name: Name of the selection
+        :type name: str
+        :param use_obs_value: If `True`, make selection on
+        observed_value. `False` by default.
+        :type use_obs_value: bool
+        :param use_distance: If `True` make selection on distance.
+        `False` by default.
+        :type use_distance: bool
+        :param use_luminosity: If `True` make selection on luminosity.
+        `False` by default.
+        :type use_luminosity: bool
+        :param use_flux: If `True` make selection on flux. `False` by default.
+        :type use_flux: bool
+        :param vmin: Minimum value of selection
+        :type vmin: :class:`SelectionParameter`
+        :param vmax: Maximum value of selection
+        :type vmax: :class:`SelectionParameter`
+        """
 
-        super(BoxSelection, self).__init__(name=name,
-                                           use_distance=use_distance,
-                                           use_luminosity=use_luminosity,
-                                           use_obs_value=use_obs_value,
-                                           use_flux=use_flux)
+        super(BoxSelection, self).__init__(
+            name=name,
+            use_distance=use_distance,
+            use_luminosity=use_luminosity,
+            use_obs_value=use_obs_value,
+            use_flux=use_flux,
+        )
 
     def draw(self, size: int) -> np.ndarray:
 
@@ -98,20 +132,43 @@ class LowerBound(SelectionProbabilty):
 
     boundary = SelectionParameter()
 
-    def __init__(self,
-                 name="Hard selection",
-                 use_obs_value: bool = False,
-                 use_distance=False,
-                 use_luminosity=False,
-                 use_flux: bool = False):
+    def __init__(
+        self,
+        name: str = "Hard selection",
+        use_obs_value: bool = False,
+        use_distance: bool = False,
+        use_luminosity: bool = False,
+        use_flux: bool = False,
+    ):
         """
-        hard selection above the boundary
+        A hard, lower bound selection on obs_value, distance,
+        luminosity or flux.
+
+        Selects values >= ``boundary``.
+
+        :param name: Name of the selection
+        :type name: str
+        :param use_obs_value: If `True`, make selection on
+        observed_value. `False` by default.
+        :type use_obs_value: bool
+        :param use_distance: If `True` make selection on distance.
+        `False` by default.
+        :type use_distance: bool
+        :param use_luminosity: If `True` make selection on luminosity.
+        `False` by default.
+        :type use_luminosity: bool
+        :param use_flux: If `True` make selection on flux. `False` by default.
+        :type use_flux: bool
+        :param boundary: Value of the selection boundary
+        :type boundary: :class:`SelectionParameter`
         """
-        super(LowerBound, self).__init__(name=name,
-                                         use_distance=use_distance,
-                                         use_luminosity=use_luminosity,
-                                         use_obs_value=use_obs_value,
-                                         use_flux=use_flux)
+        super(LowerBound, self).__init__(
+            name=name,
+            use_distance=use_distance,
+            use_luminosity=use_luminosity,
+            use_obs_value=use_obs_value,
+            use_flux=use_flux,
+        )
 
     def draw(self, size: int) -> None:
 
@@ -138,20 +195,43 @@ class UpperBound(SelectionProbabilty):
 
     boundary = SelectionParameter()
 
-    def __init__(self,
-                 name="Hard selection",
-                 use_obs_value: bool = False,
-                 use_distance=False,
-                 use_luminosity=False,
-                 use_flux: bool = False):
+    def __init__(
+        self,
+        name: str = "Hard selection",
+        use_obs_value: bool = False,
+        use_distance: bool = False,
+        use_luminosity: bool = False,
+        use_flux: bool = False,
+    ):
         """
-        hard selection below the boundary
+        A hard, upper bound selection on obs_value, distance,
+        luminosity or flux.
+
+        Selects values <= ``boundary``.
+
+        :param name: Name of the selection
+        :type name: str
+        :param use_obs_value: If `True`, make selection on
+        observed_value. `False` by default.
+        :type use_obs_value: bool
+        :param use_distance: If `True` make selection on distance.
+        `False` by default.
+        :type use_distance: bool
+        :param use_luminosity: If `True` make selection on luminosity.
+        `False` by default.
+        :type use_luminosity: bool
+        :param use_flux: If `True` make selection on flux. `False` by default.
+        :type use_flux: bool
+        :param boundary: Value of the selection boundary
+        :type boundary: :class:`SelectionParameter`
         """
-        super(UpperBound, self).__init__(name=name,
-                                         use_distance=use_distance,
-                                         use_luminosity=use_luminosity,
-                                         use_obs_value=use_obs_value,
-                                         use_flux=use_flux)
+        super(UpperBound, self).__init__(
+            name=name,
+            use_distance=use_distance,
+            use_luminosity=use_luminosity,
+            use_obs_value=use_obs_value,
+            use_flux=use_flux,
+        )
 
     def draw(self, size: int) -> None:
 
@@ -179,25 +259,45 @@ class SoftSelection(SelectionProbabilty):
     boundary = SelectionParameter()
     strength = SelectionParameter(vmin=0)
 
-    def __init__(self,
-                 name="Soft Selection",
-                 use_obs_value: bool = False,
-                 use_distance=False,
-                 use_luminosity=False,
-                 use_flux: bool = False) -> None:
+    def __init__(
+        self,
+        name: str = "Soft Selection",
+        use_obs_value: bool = False,
+        use_distance: bool = False,
+        use_luminosity: bool = False,
+        use_flux: bool = False,
+    ) -> None:
         """
-        selection using an inverse logit function either on the 
-        log are linear value of the parameter
-
-        :param boundary: center of the logit
-        :param strength: width of the logit
+        A soft selection using an inverse logit function either on the
+        log or linear value of the observed_value, distance,
+        luminosity or flux.
+
+        :param name: Name of the selection
+        :type name: str
+        :param use_obs_value: If `True`, make selection on
+        observed_value. `False` by default.
+        :type use_obs_value: bool
+        :param use_distance: If `True` make selection on distance.
+        `False` by default.
+        :type use_distance: bool
+        :param use_luminosity: If `True` make selection on luminosity.
+        `False` by default.
+        :type use_luminosity: bool
+        :param use_flux: If `True` make selection on flux. `False` by default.
+        :type use_flux: bool
+        :param boundary: Center of the inverse logit
+        :type boundary: :class:`SelectionParameter`
+        :param strength: Width of the logit
+        :type strength: :class:`SelectionParameter`
         """
 
-        super(SoftSelection, self).__init__(name=name,
-                                            use_distance=use_distance,
-                                            use_luminosity=use_luminosity,
-                                            use_obs_value=use_obs_value,
-                                            use_flux=use_flux)
+        super(SoftSelection, self).__init__(
+            name=name,
+            use_distance=use_distance,
+            use_luminosity=use_luminosity,
+            use_obs_value=use_obs_value,
+            use_flux=use_flux,
+        )
 
     def draw(self, size: int, use_log=True) -> None:
 
@@ -230,6 +330,10 @@ def draw(self, size: int, use_log=True) -> None:
                                               size=len(values)).astype(bool)
 
     __all__ = [
-        "UnitySelection", "BernoulliSelection", "BoxSelection", "UpperBound",
-        "UpperBound", "SoftSelection"
+        "UnitySelection",
+        "BernoulliSelection",
+        "BoxSelection",
+        "UpperBound",
+        "UpperBound",
+        "SoftSelection",
     ]
diff --git a/popsynth/selection_probability/selection_probability.py b/popsynth/selection_probability/selection_probability.py
index a11afcb..705fb7a 100644
--- a/popsynth/selection_probability/selection_probability.py
+++ b/popsynth/selection_probability/selection_probability.py
@@ -1,5 +1,5 @@
 import abc
-from typing import Any, Dict, Optional
+from typing import Dict
 
 import numpy as np
 from class_registry import AutoRegister
@@ -18,12 +18,31 @@ class SelectionParameter(Parameter):
 class SelectionProbabilty(object,
                           metaclass=AutoRegister(selection_registry,
                                                  base_type=ParameterMeta)):
-    def __init__(self,
-                 name: str = "name",
-                 use_obs_value: bool = False,
-                 use_distance: bool = False,
-                 use_luminosity: bool = False,
-                 use_flux: bool = False) -> None:
+    def __init__(
+        self,
+        name: str = "name",
+        use_obs_value: bool = False,
+        use_distance: bool = False,
+        use_luminosity: bool = False,
+        use_flux: bool = False,
+    ) -> None:
+        """
+        Base class for selections on a population.
+
+        :param name: Name of the selection
+        :type name: str
+        :param use_obs_value: If `True`, make selection on
+        observed_value. `False` by default.
+        :type use_obs_value: bool
+        :param use_distance: If `True` make selection on distance.
+        `False` by default.
+        :type use_distance: bool
+        :param use_luminosity: If `True` make selection on luminosity.
+        `False` by default.
+        :type use_luminosity: bool
+        :param use_flux: If `True` make selection on flux. `False` by default.
+        :type use_flux: bool
+        """
 
         self._name = name  # type: str
 
@@ -57,41 +76,30 @@ def __add__(self, other):
 
     def set_luminosity(self, luminosity: np.ndarray) -> None:
         """
-        set the luminosity of the selection
-
-        :param luminosity:
-        :returns:
-        :rtype:
-
+        Set the luminosity of the selection.
         """
 
         self._luminosity = luminosity  # type: np.ndarray
 
     def set_distance(self, distance: np.ndarray) -> None:
         """
-        set the distance of the selection
-
-        :param distance:
-        :returns:
-        :rtype:
-
+        Set the distance of the selection.
         """
 
         self._distance = distance  # type: np.ndarray
 
     def set_observed_flux(self, observed_flux: np.ndarray) -> None:
         """
-        set the observed flux of the selection
-
-        :param luminosity:
-        :returns:
-        :rtype:
-
+        Set the observed flux of the selection.
         """
 
         self._observed_flux = observed_flux  # type: np.ndarray
 
     def set_observed_value(self, observed_value: np.ndarray) -> None:
+        """
+        Set the observed value of the selection.
+        """
+
         self._observed_value = observed_value  # type: np.ndarray
 
     @abc.abstractclassmethod
@@ -101,7 +109,7 @@ def draw(self, size: int) -> None:
 
     def reset(self):
         """
-        resest the selector
+        Reset the selector.
         """
         if self._is_sampled:
 
@@ -177,6 +185,9 @@ class DummySelection(SelectionProbabilty):
     _selection_name = "DummySelection"
 
     def __init__(self):
+        """
+        A Dummy selection for testing.
+        """
 
         super(DummySelection, self).__init__(name="dummy")
 
diff --git a/popsynth/selection_probability/spatial_selection.py b/popsynth/selection_probability/spatial_selection.py
index 7e62381..18ffc15 100644
--- a/popsynth/selection_probability/spatial_selection.py
+++ b/popsynth/selection_probability/spatial_selection.py
@@ -1,13 +1,20 @@
-from popsynth.distribution import SpatialDistribution
+from astropy.coordinates import SkyCoord
 
-from .selection_probability import SelectionProbabilty
+from popsynth.distribution import SpatialDistribution
+from .selection_probability import SelectionProbabilty, SelectionParameter
 
 
 class SpatialSelection(SelectionProbabilty):
 
     _selection_name = "SpatialSelection"
 
-    def __init__(self, name) -> None:
+    def __init__(self, name: str) -> None:
+        """
+        A generic spatial selection.
+
+        :param name: Name of the selection
+        :type name: str
+        """
 
         super(SpatialSelection, self).__init__(name)
 
@@ -16,11 +23,46 @@ def __init__(self, name) -> None:
     def set_spatial_distribution(
             self, spatial_distribtuion: SpatialDistribution) -> None:
         """
-        set the spatial distribution
+        Set the spatial distribution for the selection.
+
+        :param spatial_distribution: The spatial_distribution
+        :type spatial_distribution: :class:`SpatialDistribution`
         """
         assert isinstance(spatial_distribtuion, SpatialDistribution)
 
         self._spatial_distribution: SpatialDistribution = spatial_distribtuion
 
 
-__all__ = ["SpatialSelection"]
+class GalacticPlaneSelection(SpatialSelection):
+
+    _selection_name = "GalacticPlaneSelection"
+
+    b_limit = SelectionParameter(vmin=0, vmax=90)
+
+    def __init__(self, name: str = "galactic plane selector"):
+        """
+        A selection that excludes objects near the galactic plane.
+
+        :param name: Name of the selection
+        :type name: str
+        :param b_limit: Limit around Galactic plane to exclude in
+        Galactic latitude and in units of degrees
+        :type b_limit: :class:`SelectionParameter`
+        """
+        super(GalacticPlaneSelection, self).__init__(name=name)
+
+    def draw(self, size: int):
+
+        c = SkyCoord(
+            self._spatial_distribution.ra,
+            self._spatial_distribution.dec,
+            unit="deg",
+            frame="icrs",
+        )
+
+        b = c.galactic.b.deg
+
+        self._selection = (b >= self.b_limit) | (b <= -self.b_limit)
+
+
+__all__ = ["SpatialSelection", "GalacticPlaneSelection"]
diff --git a/popsynth/utils/array_to_cmap.py b/popsynth/utils/array_to_cmap.py
index ba8e67d..b68d0c0 100644
--- a/popsynth/utils/array_to_cmap.py
+++ b/popsynth/utils/array_to_cmap.py
@@ -2,16 +2,18 @@
 import matplotlib.pyplot as plt
 
 
-def array_to_cmap(values, cmap, use_log=False):
+def array_to_cmap(values, cmap, use_log: bool = False):
     """
     Generates a color map and color list that is normalized
     to the values in an array. Allows for adding a 3rd dimension
-    onto a plot
-
-    :param values: a list a values to map into a cmap
-    :param cmap: the mpl colormap to use
-    :param use_log: if the mapping should be done in log space
-
+    onto a plot.
+
+    :param values: A list a values to map into a cmap
+    :param cmap: The mpl colormap to use
+    :param use_log: `True` if the mapping should be done in log space
+    `False` by default.
+    :type use_log: bool
+    :returns: A color map and a normalized color list
     """
 
     if use_log:
diff --git a/popsynth/utils/configuration.py b/popsynth/utils/configuration.py
index 63648d7..aecc1e0 100644
--- a/popsynth/utils/configuration.py
+++ b/popsynth/utils/configuration.py
@@ -1,18 +1,19 @@
-import os
-import warnings
 from pathlib import Path
 from dataclasses import dataclass
 
 from omegaconf import OmegaConf
 
+# Path to configuration
+
 _config_path = Path("~/.config/popsynth/").expanduser()
 
 _config_name = Path("popsynth_config.yml")
 
 _config_file = _config_path / _config_name
 
+# Define structure of configuration with dataclasses
+
 
-# Define structure with dataclasses
 @dataclass
 class LogConsole:
 
@@ -53,7 +54,7 @@ class PopSynthConfig:
 # Read the default config
 popsynth_config: PopSynthConfig = OmegaConf.structured(PopSynthConfig)
 
-# Merge with local config
+# Merge with local config if it exists
 if _config_file.is_file():
 
     _local_config = OmegaConf.load(_config_file)
@@ -61,7 +62,7 @@ class PopSynthConfig:
     popsynth_config: PopSynthConfig = OmegaConf.merge(popsynth_config,
                                                       _local_config)
 
-# Write defaults
+# Write defaults if not
 else:
 
     # Make directory if needed
diff --git a/popsynth/utils/cosmology.py b/popsynth/utils/cosmology.py
index cd6f151..5fba5f4 100644
--- a/popsynth/utils/cosmology.py
+++ b/popsynth/utils/cosmology.py
@@ -3,7 +3,7 @@
 import numpy as np
 from astropy.constants import c
 from astropy.cosmology import WMAP9 as cosmo
-from numba import bool_, float32, float64
+from numba import float32
 from numba.experimental import jitclass
 
 from popsynth.utils.configuration import popsynth_config
@@ -12,8 +12,8 @@
 
 spec = [
     #    ("_is_setup", bool_),
-    ('Om', float32),
-    ('h0', float32),
+    ("Om", float32),
+    ("h0", float32),
     # ('_dh', float64),
     # ('_om_reduced', float64),
     # ('_om_sqrt', float64),
@@ -26,9 +26,20 @@
 
 @jitclass(spec)
 class Cosmology(object):
-    def __init__(self,
-                 Om=popsynth_config["cosmology"]["Om"],
-                 h0=popsynth_config["cosmology"]["h0"]):
+    def __init__(
+        self,
+        Om: float = popsynth_config["cosmology"]["Om"],
+        h0: float = popsynth_config["cosmology"]["h0"],
+    ):
+        """
+        Cosmological model used in relevant calculations.
+        Flat LambdaCDM. Can be set using configuration.
+
+        :param Om: Omega matter
+        :type Om: float
+        :param h0: Hubble constant in km s^-1 Mpc^-1
+        :type h0: float
+        """
 
         #      self._is_setup = False
 
@@ -79,7 +90,9 @@ def xx(self, z):
 
     def luminosity_distance(self, z):
         """
-        dL in cm
+        Luminosity distance in units of cm.
+
+        :param z: Redshift
         """
         x = self.xx(z)
         z1 = 1.0 + z
@@ -98,7 +111,9 @@ def comoving_transverse_distance(self, z):
 
     def differential_comoving_volume(self, z):
         """
-        differential comoving volume in Gpc2
+        Differential comoving volume in Gpc^3 sr^-1
+
+        :param z: Redshift
         """
 
         td = self.comoving_transverse_distance(z) / 3.086e24
diff --git a/popsynth/utils/hdf5_utils.py b/popsynth/utils/hdf5_utils.py
index 6c0b51d..90407b8 100644
--- a/popsynth/utils/hdf5_utils.py
+++ b/popsynth/utils/hdf5_utils.py
@@ -1,13 +1,16 @@
-import os
-
 import h5py
 import numpy as np
+from typing import Dict, Any
 
 
-def recursively_save_dict_contents_to_group(h5file, path, dic):
+def recursively_save_dict_contents_to_group(h5file, path, dic: Dict[Any, Any]):
     """
-    load files from hdf5
+    Save dict to HDF5.
 
+    :param h5file: HDF5 file
+    :param path: Path inside file
+    :param dic: Dictionary to save
+    :type dic: Dict[Any, Any]
     """
     for key, item in dic.items():
 
@@ -62,7 +65,12 @@ def recursively_save_dict_contents_to_group(h5file, path, dic):
 
 def recursively_load_dict_contents_from_group(h5file, path):
     """
-    load files from hdf5
+    Load files from hdf5.
+
+    :param h5file: HDF5 file
+    :param path: Path in file
+    :returns: A dictionary
+    :rtype: Dict[Any, Any]
     """
     ans = {}
 
@@ -89,12 +97,11 @@ def recursively_load_dict_contents_from_group(h5file, path):
     return ans
 
 
-# these two functions add and clean
-# a networkx graph dict so that it will be stored in
-# an hdf5 file
-
-
 def fill_graph_dict(graph_dict, fill_value=1):
+    """
+    Fill a networkx graph dict so that it
+    can be stored in an HDF5 file.
+    """
     new_dict = {}
 
     for k, v in graph_dict.items():
@@ -108,6 +115,10 @@ def fill_graph_dict(graph_dict, fill_value=1):
 
 
 def clean_graph_dict(graph_dict):
+    """
+    Clean networkx graph dict so that it
+    can be stored in an HDF5 file.
+    """
     new_dict = {}
 
     for k, v in graph_dict.items():
diff --git a/popsynth/utils/logging.py b/popsynth/utils/logging.py
index 2873909..bcbc080 100644
--- a/popsynth/utils/logging.py
+++ b/popsynth/utils/logging.py
@@ -8,8 +8,7 @@
 from colorama import Back, Fore, Style
 
 from popsynth.utils.configuration import popsynth_config
-from popsynth.utils.package_data import (get_path_of_log_dir,
-                                         get_path_of_log_file)
+from popsynth.utils.package_data import get_path_of_log_dir, get_path_of_log_file
 
 colorama.deinit()
 colorama.init(strip=False)
@@ -105,7 +104,7 @@ def __init__(
         popsynth_console_log_handler,
     ):
         """
-        A container to store the stat of the logs
+        A container to store the state of the logs.
         """
 
         # attach the log handlers
@@ -166,7 +165,7 @@ def debug_logs(self):
 
 def silence_warnings():
     """
-    supress warning messages in console and file usr logs
+    Supress warning messages in console and file usr logs.
     """
 
     popsynth_usr_log_handler.addFilter(warning_filter)
@@ -175,7 +174,7 @@ def silence_warnings():
 
 def activate_warnings():
     """
-    supress warning messages in console and file usr logs
+    Supress warning messages in console and file usr logs.
     """
 
     popsynth_usr_log_handler.removeFilter(warning_filter)
@@ -189,7 +188,7 @@ def update_logging_level(level):
 
 def silence_logs():
     """
-    Turn off all logging
+    Turn off all logging.
     """
 
     # handle dev logs independently
@@ -210,7 +209,7 @@ def silence_progress_bars():
 
 def quiet_mode():
     """
-    turn off all logging and progress bars
+    Turn off all logging and progress bars.
     """
 
     silence_progress_bars()
@@ -221,7 +220,7 @@ def quiet_mode():
 
 def loud_mode():
     """
-    turn on all progress bars and logging
+    Turn on all progress bars and logging.
     """
 
     show_progress_bars()
@@ -232,7 +231,7 @@ def loud_mode():
 
 def activate_logs():
     """
-    re-activate silenced logs
+    Re-activate silenced logs.
     """
 
     # handle dev logs independently
@@ -243,7 +242,7 @@ def activate_logs():
 
 def debug_mode():
     """
-    activate debug in the console
+    Activate debug in the console.
     """
 
     # store state and switch console to debug
@@ -269,6 +268,11 @@ def silence_console_log():
 
 
 def setup_logger(name):
+    """
+    Set up a new logger.
+
+    :param name: Name of the logger
+    """
 
     # A logger with name name will be created
     # and then add it to the print stream
diff --git a/popsynth/utils/meta.py b/popsynth/utils/meta.py
index 73fdc8c..6f2b9da 100644
--- a/popsynth/utils/meta.py
+++ b/popsynth/utils/meta.py
@@ -6,11 +6,24 @@
 
 
 class Parameter(object):
-    def __init__(self,
-                 default: Optional[float] = None,
-                 vmin: Optional[float] = None,
-                 vmax: Optional[float] = None,
-                 free: bool = True):
+    def __init__(
+        self,
+        default: Optional[float] = None,
+        vmin: Optional[float] = None,
+        vmax: Optional[float] = None,
+        free: bool = True,
+    ):
+        """
+        Parameter base class.
+
+        :param default: Default parameter value
+        :type default: Optional[float]
+        :param vmin: Minimum parameter value
+        :type vmin: Optional[float]
+        :param vmax: Maximum parameter value
+        :type vmax: Optional[float]
+        :param free: If `True`, parameter is free
+        """
 
         self.name = None  # type: Union[None, str]
         self._vmin = vmin  # type: Optional[float]
@@ -30,8 +43,8 @@ def __get__(self, obj, type=None) -> object:
 
                 obj._parameter_storage[self.name] = self._default
 
-            assert obj._parameter_storage[
-                self.name] is not None, "parameters must have values!"
+            assert (obj._parameter_storage[self.name]
+                    is not None), "parameters must have values!"
 
             return obj._parameter_storage[self.name]
 
@@ -39,8 +52,8 @@ def __get__(self, obj, type=None) -> object:
 
             obj._parameter_storage[self.name] = self._default
 
-        assert obj._parameter_storage[
-            self.name] is not None, "parameters must have values!"
+        assert (obj._parameter_storage[self.name]
+                is not None), "parameters must have values!"
 
         return obj._parameter_storage[self.name]
 
@@ -49,7 +62,7 @@ def __set__(self, obj, value) -> None:
 
         if not self._free:
 
-            log.error(f"{self.name} is fixed and cannot me set")
+            log.error(f"{self.name} is fixed and cannot be set")
             raise RuntimeError()
 
         if self._vmin is not None:
diff --git a/popsynth/utils/package_data.py b/popsynth/utils/package_data.py
index 6f7ac69..9d8af65 100644
--- a/popsynth/utils/package_data.py
+++ b/popsynth/utils/package_data.py
@@ -1,8 +1,6 @@
 from pathlib import Path
 import pkg_resources
 
-from popsynth.utils.configuration import popsynth_config
-
 
 def get_path_of_data_file(data_file) -> Path:
     file_path = pkg_resources.resource_filename("popsynth",
diff --git a/popsynth/utils/registry.py b/popsynth/utils/registry.py
index ffa6471..e5722a0 100644
--- a/popsynth/utils/registry.py
+++ b/popsynth/utils/registry.py
@@ -1,4 +1,4 @@
-from class_registry import AutoRegister, ClassRegistry
+from class_registry import ClassRegistry
 
 auxiliary_parameter_registry = ClassRegistry("_auxiliary_sampler_name")
 
@@ -26,6 +26,7 @@ def list_available_auxiliary_samplers():
 
 
 __all__ = [
-    "auxiliary_parameter_registry", "distribution_registry",
-    "distribution_registry"
+    "auxiliary_parameter_registry",
+    "distribution_registry",
+    "distribution_registry",
 ]
diff --git a/popsynth/utils/rejection_sample.py b/popsynth/utils/rejection_sample.py
index cacd325..49e82e2 100644
--- a/popsynth/utils/rejection_sample.py
+++ b/popsynth/utils/rejection_sample.py
@@ -4,15 +4,13 @@
 
 # @jit(parallel=False, forceobj=True)
 def rejection_sample(size, ymax, xmax, func):
-    """FIXME! briefly describe function
-
-    :param size:
-    :param ymax:
-    :param xmax:
-    :param func:
-    :returns:
-    :rtype:
+    """
+    Rejection sample ``func`` up to ``ymax`` and ``xmax``.
 
+    :param size: Number of samples
+    :param ymax: Maximum value of y
+    :param xmax: Maximum value of x
+    :param func: Function
     """
 
     r_out = np.zeros(size)
diff --git a/popsynth/utils/spherical_geometry.py b/popsynth/utils/spherical_geometry.py
index 3745a87..5b00427 100644
--- a/popsynth/utils/spherical_geometry.py
+++ b/popsynth/utils/spherical_geometry.py
@@ -2,6 +2,10 @@
 
 
 def sample_theta_phi(size: int):
+    """
+    Sample ``size`` samples uniformly on the
+    surface of the unit sphere.
+    """
 
     theta = np.arccos(1 - 2 * np.random.uniform(0.0, 1.0, size=size))
     phi = np.random.uniform(0, 2 * np.pi, size=size)
@@ -10,6 +14,9 @@ def sample_theta_phi(size: int):
 
 
 def xyz(r, theta, phi):
+    """
+    Convert spherical coordinates to Cartesian.
+    """
     x = r * np.sin(theta) * np.cos(phi)
     y = r * np.sin(theta) * np.sin(phi)
     z = r * np.cos(theta)


Test patch:
diff --git a/popsynth/test/test_selections.py b/popsynth/test/test_selections.py
index 5cb3a40..1292686 100644
--- a/popsynth/test/test_selections.py
+++ b/popsynth/test/test_selections.py
@@ -1,32 +1,6 @@
-import numpy as np
-from astropy.coordinates import SkyCoord
-
 import popsynth
 
-
-class GalacticPlaneSelection(popsynth.SpatialSelection):
-
-    _selection_name = "GalacticPlaneSelection"
-
-    b_limit = popsynth.SelectionParameter(vmin=0, vmax=90)
-
-    def __init__(self, name="mw plane selector"):
-        """
-        places a limit above the galactic plane for objects
-        """
-        super(GalacticPlaneSelection, self).__init__(name=name)
-
-    def draw(self, size: int):
-
-        g_coor = SkyCoord(
-            self._spatial_distribution.ra,
-            self._spatial_distribution.dec,
-            unit="deg",
-            frame="icrs",
-        ).transform_to("galactic")
-
-        self._selection = (g_coor.b.deg >= self.b_limit) | (g_coor.b.deg <=
-                                                            -self.b_limit)
+from popsynth.selection_probability.spatial_selection import GalacticPlaneSelection
 
 
 def test_spatial_selection():