fix the broken docs

CHANGELOG.md

@@ -1,5 +1,11 @@
 # All notable changes will be documented in this file
 
+## [0.2.1] 2023-11-24
+Version 0.2.1 release of `redback_surrogates`
+
+### Changed
+- Minor changes to the documentation/interface for afterglow emulator.
+
 ## [0.2] 2023-11-24
 Version 0.2 release of `redback_surrogates`
 

docs/conf.py

@@ -24,6 +24,7 @@
 # The full version, including alpha/beta/rc tags
 release = 'alpha'
 autosummary_generate = True
+autodoc_mock_imports = ["pickle", "scipy", "kilonovanet"]
 
 source_suffix = ['.rst', '.md', '.txt', '.ipynb', '.py']
 # -- General configuration ---------------------------------------------------

redback_surrogates/afterglowmodels.py

@@ -2,6 +2,7 @@
 from scipy import interpolate
 import numpy as np
 import os
+from redback_surrogates.utils import citation_wrapper
 dirname = os.path.dirname(__file__)
 
 with open(f"{dirname}/surrogate_data/onax_redback.pkl", "rb") as f_on:
@@ -20,14 +21,15 @@
 
 def _shape_data(thv, loge0, thc, logn0, p, logepse, logepsb, g0,frequency):
     if isinstance(frequency, (int, float)) == True:
-        test_data= np.array([np.log10(thv) , loge0 , np.log10(thc), logn0, p, logepse, logepsb, np.log10(g0), frequency]).reshape(1,-1)
+        test_data = np.array([np.log10(thv) , loge0 , np.log10(thc), logn0, p, logepse, logepsb, np.log10(g0), frequency]).reshape(1,-1)
     else:
-        test_data= []
+        test_data = []
         for f in frequency:
             test_data.append([np.log10(thv) , loge0 , np.log10(thc), logn0, p, logepse, logepsb, np.log10(g0), f])
     return np.array(test_data)    
-
-def tophat_emulator(new_time, thv, loge0, thc, logn0, p, logepse, logepsb, g0, frequency):
+
+@citation_wrapper("Wallace and Sarin in prep.")
+def tophat_emulator(new_time, thv, loge0, thc, logn0, p, logepse, logepsb, g0, **kwargs):
     """
     tophat afterglow model using trained mpl regressor
 
@@ -40,24 +42,26 @@ def tophat_emulator(new_time, thv, loge0, thc, logn0, p, logepse, logepsb, g0, f
     :param logepse: log10 fraction of thermal energy in electrons
     :param logepsb: log10 fraction of thermal energy in magnetic field
     :param g0: initial lorentz factor
+    :param kwargs: extra arguments for the model
     :param frequency: frequency of the band to view in- single number or same length as time array
     :return: flux density at each time for given frequency
     """
-
+
+    frequency = kwargs['frequency']
     test_data = _shape_data(thv, loge0, thc, logn0, p, logepse, logepsb, g0,frequency)
-    logtime= np.logspace(2.94,7.41,100)/86400
-    
-    if thv<=thc:
-        xtests= scalerx_on.transform(test_data)
-        prediction= model_on.predict(xtests)
+    logtime = np.logspace(2.94,7.41,100)/86400
+
+    if thv <= thc:
+        xtests = scalerx_on.transform(test_data)
+        prediction = model_on.predict(xtests)
         prediction = np.exp(scalery_on.inverse_transform(prediction))
-    elif thv>thc:
-        xtests= scalerx_off.transform(test_data)
-        prediction= model_off.predict(xtests)
+    else:
+        xtests = scalerx_off.transform(test_data)
+        prediction = model_off.predict(xtests)
         prediction = np.exp(scalery_off.inverse_transform(prediction))
 
     afterglow = interpolate.interp1d(logtime, prediction, kind='linear', fill_value='extrapolate')
-    fluxd= afterglow(new_time)
+    fluxd = afterglow(new_time)
 
     if test_data.shape == (1,9):
         return fluxd[0]

setup.py

@@ -4,7 +4,7 @@
     long_description = fh.read()
 
 setup(name='redback_surrogates',
-      version='0.2',
+      version='0.2.1',
       description='A surrogates package to work with redback, the bayesian inference package for electromagnetic transients',
       long_description=long_description,
       long_description_content_type="text/markdown",