Merge pull request #81 from mhvk/move-modeling-stuff-to-new-module

Move modeling stuff to new module

docs/index.rst

@@ -80,4 +80,5 @@ Reference/API
 .. automodapi:: screens.dynspec
 .. automodapi:: screens.conjspec
 .. automodapi:: screens.remap
+.. automodapi:: screens.modeling
 .. automodapi:: screens.visualization

examples/brisken.py

@@ -2,15 +2,18 @@
 """Fit part of the Brisken dynamic spectrum.
 
 Example runs on Brisken file 'dynamic_spectrum_arar.npz'
+# On CITA machines can use
+# ln -s /mnt/scratch-lustre/simard/GB057/dynamic_spectra/dynamic_spectrum_arar.npz
 """
 
 import numpy as np
 from astropy import units as u
 from matplotlib import pyplot as plt
 from astropy.visualization import quantity_support
 
-from screens.visualization import ThetaTheta
 from screens import DynamicSpectrum
+from screens.modeling import DynamicSpectrumModel
+from screens.visualization import ThetaTheta
 
 
 quantity_support()
@@ -22,9 +25,10 @@
 sel_t, sel_f = slice(292, 505), slice(50000, 50300)
 t = t_full[sel_t]
 f = f_full[sel_f]
-ds = DynamicSpectrum(a['I'][sel_t, sel_f], t=t[:, np.newaxis],
-                     f=f, noise=1, d_eff=1.4*u.kpc,
-                     mu_eff=50*u.mas/u.yr)
+dynspec = DynamicSpectrum(a['I'][sel_t, sel_f], t=t[:, np.newaxis],
+                          f=f, noise=1)
+ds = DynamicSpectrumModel(dynspec, d_eff=1.4*u.kpc,
+                          mu_eff=50*u.mas/u.yr)
 
 
 tau_max = 350*u.us

examples/dm_gradient.py

@@ -14,7 +14,7 @@
 
 from baseband_tasks.dm import DispersionMeasure
 
-from screens.fields import dynamic_field, phasor
+from screens.fields import dynamic_field
 from screens.dynspec import DynamicSpectrum as DS
 from screens.conjspec import ConjugateSpectrum as CS
 
@@ -131,13 +131,13 @@
                extent=ss_extent, cmap='Greys', vmin=-9, vmax=-2)
     plt.xlabel(rf"$f_{{D}}\ ({nut.fd.unit.to_string('latex')[1:-1]})$")
     plt.ylabel(rf"$\tau\ ({nut.tau.unit.to_string('latex')[1:-1]})$")
-    plt.title(rf"$\nu - \nu t$")
+    plt.title(r"$\nu - \nu t$")
     plt.colorbar()
 
     plt.subplot(3, 3, 6+1+i)
     plt.imshow(np.log10(nut2.secspec.T), origin='lower', aspect='auto',
                extent=ss_extent, cmap='Greys', vmin=-9, vmax=-2)
     plt.xlabel(rf"$f_{{D}}\ ({nut2.fd.unit.to_string('latex')[1:-1]})$")
     plt.ylabel(rf"$\tau\ ({nut2.tau.unit.to_string('latex')[1:-1]})$")
-    plt.title(rf"$\nu - \nu t_{{corr}}$")
+    plt.title(r"$\nu - \nu t_{corr}$")
     plt.colorbar()

examples/screen2ds.py

@@ -120,8 +120,8 @@
 
 tau_max = (1./(f[3]-f[0])).to(u.us)
 th_g = theta_grid(d_eff, mu_eff, fobs=fobs,
-                  dtau=1/f.ptp(), tau_max=tau_max,
-                  dfd=1/t.ptp(), fd_max=1*u.Hz)
+                  dtau=1/np.ptp(f), tau_max=tau_max,
+                  dfd=1/np.ptp(t), fd_max=1*u.Hz)
 fd_g = (d_eff/const.c*mu_eff*fobs*th_g).to(
     u.mHz, equivalencies=u.dimensionless_angles())
 tau_g = (d_eff/(2*const.c)*th_g**2).to(

examples/theta.py

@@ -14,6 +14,7 @@
 
 from screens.fields import dynamic_field, theta_grid
 from screens.dynspec import DynamicSpectrum
+from screens.modeling import DynamicSpectrumModel
 from screens.visualization import ThetaTheta
 
 
@@ -51,12 +52,13 @@
 # Assue we're not too close to max tau in secondary spectrum.
 tau_max = (1./(f[3]-f[0])).to(u.us)
 th_r = theta_grid(d_eff, mu_eff, fobs=fobs,
-                  dtau=1/f.ptp(), tau_max=tau_max,
-                  dfd=1/t.ptp(), fd_max=1*u.Hz)
+                  dtau=1/np.ptp(f), tau_max=tau_max,
+                  dfd=1/np.ptp(t), fd_max=1*u.Hz)
 
 # Create a DynamicSpectrum just to use its `theta_theta` method.
-ds = DynamicSpectrum(dynspec, f=f, t=t, noise=noise, d_eff=d_eff,
-                     mu_eff=mu_eff, theta=th_r)
+ds = DynamicSpectrumModel(
+    DynamicSpectrum(dynspec, f=f, t=t, noise=noise),
+    d_eff=d_eff, mu_eff=mu_eff, theta=th_r)
 
 th_th = ds.theta_theta()
 
@@ -71,7 +73,7 @@
 
 # Calculate eigenvectors for inferred theta-theta.
 
-w, v = scipy.linalg.eigh(th_th, eigvals=(th_r.size-1, th_r.size-1))
+w, v = scipy.linalg.eigh(th_th, subset_by_index=(th_r.size-1, th_r.size-1))
 
 # Ideally, the eigenvalue is 1, but we want a normalized solution anyway,
 # so just use properly normalized eigenvector.

examples/try_dyn_fit.py

@@ -12,14 +12,15 @@
 from matplotlib import pyplot as plt
 from astropy.visualization import quantity_support
 
-from screens.visualization import ThetaTheta
 from screens import DynamicSpectrum
+from screens.modeling import DynamicSpectrumModel
+from screens.visualization import ThetaTheta
 
 
 quantity_support()
 
-dyn_chi2 = DynamicSpectrum.fromfile('dynspec.h5', d_eff=1.*u.kpc,
-                                    mu_eff=100*u.mas/u.yr)
+ds = DynamicSpectrum.fromfile('dynspec.h5')
+dyn_chi2 = DynamicSpectrumModel(ds, d_eff=1.*u.kpc, mu_eff=100*u.mas/u.yr)
 dyn_chi2.theta = dyn_chi2.theta_grid(
     tau_max=(1./(dyn_chi2.f[3]-dyn_chi2.f[0])).to(u.us))
 dyn_chi2.locate_mu_eff(np.arange(98, 103) << u.mas/u.yr)

examples/try_sec_fit.py

@@ -13,21 +13,23 @@
 from astropy.visualization import quantity_support
 from scipy.linalg import eigh
 
-from screens.visualization import ThetaTheta
 from screens import DynamicSpectrum, ConjugateSpectrum
+from screens.modeling import DynamicSpectrumModel, ConjugateSpectrumModel
+from screens.visualization import ThetaTheta
 
 
 quantity_support()
 
-dyn_spec = DynamicSpectrum.fromfile('dynspec.h5', d_eff=1.*u.kpc,
-                                    mu_eff=100*u.mas/u.yr)
-conj_spec = ConjugateSpectrum.from_dynamic_spectrum(dyn_spec)
+ds = DynamicSpectrum.fromfile('dynspec.h5')
+dyn_spec = DynamicSpectrumModel(ds, d_eff=1.*u.kpc, mu_eff=100*u.mas/u.yr)
+cs = ConjugateSpectrum.from_dynamic_spectrum(ds)
+conj_spec = ConjugateSpectrumModel(cs, d_eff=1.*u.kpc, mu_eff=100*u.mas/u.yr)
 
 sec_kwargs = dict(extent=(conj_spec.fd[0, 0].value, conj_spec.fd[-1, 0].value,
                           conj_spec.tau[0].value, conj_spec.tau[-1].value),
                   cmap='Greys', vmin=-7, vmax=0, origin='lower', aspect='auto')
 plt.subplot(321)
-plt.imshow(np.log10(conj_spec.secspec).T, **sec_kwargs)
+plt.imshow(np.log10(cs.secspec).T, **sec_kwargs)
 
 conserve = True
 
@@ -45,7 +47,7 @@
 conjspec = conj_spec.conjspec.copy()
 conjspec[(conj_spec.fd == 0) | (conj_spec.tau == 0)] = 0
 # try recoving just plain power
-w_a, v_a = eigh(np.abs(th_th)**2, eigvals=(conj_spec.theta.size-1,)*2)
+w_a, v_a = eigh(np.abs(th_th)**2, subset_by_index=(conj_spec.theta.size-1,)*2)
 rec_a = v_a[:, -1] * np.sqrt(w_a[-1])
 th_th_rp = rec_a[:, np.newaxis] * rec_a
 ax = plt.subplot(324, projection=th_th_proj)
@@ -60,7 +62,7 @@
 plt.imshow(np.log10(np.maximum(sec_rp, 1e-30)).T, **sec_kwargs)
 
 # try recovering phases as well.
-w, v = eigh(th_th, eigvals=(conj_spec.theta.size-1,)*2)
+w, v = eigh(th_th, subset_by_index=(conj_spec.theta.size-1,)*2)
 recovered = v[:, -1] * np.sqrt(w[-1])
 th_th_r = recovered[:, np.newaxis] * recovered
 ax = plt.subplot(326, projection=th_th_proj)