Merge pull request #148 from abhisrkckl/pulsar

`pyvela` script

.gitignore

@@ -188,4 +188,8 @@ cython_debug/
 
 profile.pb.gz
 
-*_results.pdf
+*_results.pdf
+
+pyvela/tests/_*_out/
+pyvela_results
+__prior.json

CHANGELOG.md

@@ -3,6 +3,8 @@
 - `time_residuals`, `scaled_toa_uncertainties`, and `model_dm`, `from_pint` methods in `SPNTA`
 - NE_SW derivatives in `SolarWind`
 - `Pulsar` class
+- `pyvela` script
+- `get_free_param_units` function
 ## Changed
 ## Fixed
 ## Removed

pyproject.toml

@@ -15,7 +15,6 @@ maintainers = [
     {name = "Abhimanyu Susobhanan", email = "[email protected]"}
 ]
 
-
 requires-python = ">=3.10"
 dependencies = [
     "numpy>=2.1",
@@ -44,6 +43,9 @@ content-type = "text/x-rst"
 Homepage = "https://github.com/abhisrkckl/Vela.jl"
 Documentation = "https://abhisrkckl.github.io/Vela.jl/"
 
+[project.scripts]
+pyvela = "pyvela.pyvela_script:main"
+
 [tool.setuptools]
 zip-safe = false
 package-dir = {"" = "pyvela"}

pyvela/pyvela/pyvela_script.py

@@ -0,0 +1,112 @@
+import datetime
+import getpass
+import json
+import os
+import platform
+import shutil
+import sys
+from argparse import ArgumentParser
+
+import emcee
+import numpy as np
+import pint
+
+import pyvela
+from pyvela import SPNTA
+from pyvela import Vela as vl
+
+
+def info_dict(args):
+    return {
+        "input": {
+            "par_file": args.par_file,
+            "tim_file": args.tim_file,
+            "prior_file": args.prior_file,
+            "cheat_prior_scale": args.cheat_prior_scale,
+        },
+        "sampler": {
+            "nsteps": args.nsteps,
+            "burnin": args.burnin,
+            "thin": args.thin,
+        },
+        "env": {
+            "launch_time": datetime.datetime.now().isoformat(),
+            "user": getpass.getuser(),
+            "host": platform.node(),
+            "os": platform.platform(),
+            "julia_threads": vl.nthreads(),
+            "python": sys.version,
+            "julia": str(vl.VERSION),
+            "pyvela": pyvela.__version__,
+            "pint": pint.__version__,
+            "emcee": emcee.__version__,
+        },
+    }
+
+
+def parse_args(argv):
+    parser = ArgumentParser()
+    parser.add_argument("par_file")
+    parser.add_argument("tim_file")
+    parser.add_argument("-P", "--prior_file")
+    parser.add_argument("--cheat_prior_scale", default=10, type=float)
+    parser.add_argument("-o", "--outdir", default="pyvela_results")
+    parser.add_argument("-N", "--nsteps", default=6000, type=int)
+    parser.add_argument("-b", "--burnin", default=1500, type=int)
+    parser.add_argument("-t", "--thin", default=100, type=int)
+
+    return parser.parse_args(argv)
+
+
+def prepare_outdir(args):
+    summary_info = info_dict(args)
+
+    if not os.path.isdir(args.outdir):
+        os.mkdir(args.outdir)
+    shutil.copy(args.par_file, args.outdir)
+    shutil.copy(args.tim_file, args.outdir)
+    if args.prior_file is not None:
+        shutil.copy(args.prior_file, args.outdir)
+    with open(f"{args.outdir}/summary.json", "w") as summary_file:
+        json.dump(summary_info, summary_file, indent=4)
+
+
+def save_spnta_attrs(spnta: SPNTA, args):
+    np.savetxt(f"{args.outdir}/param_names.txt", spnta.param_names, fmt="%s")
+    np.savetxt(f"{args.outdir}/param_units.txt", spnta.param_units, fmt="%s")
+    np.savetxt(f"{args.outdir}/param_scale_factors.txt", spnta.scale_factors, fmt="%s")
+
+
+def main(argv=None):
+    args = parse_args(argv)
+    prepare_outdir(args)
+
+    spnta = SPNTA(
+        args.par_file,
+        args.tim_file,
+        cheat_prior_scale=args.cheat_prior_scale,
+        custom_priors=(args.prior_file if args.prior_file is not None else {}),
+    )
+
+    save_spnta_attrs(spnta, args)
+
+    nwalkers = spnta.ndim * 5
+    p0 = np.array(
+        [spnta.prior_transform(cube) for cube in np.random.rand(nwalkers, spnta.ndim)]
+    )
+
+    sampler = emcee.EnsembleSampler(
+        nwalkers,
+        spnta.ndim,
+        spnta.lnpost_vectorized,
+        # moves=[emcee.moves.StretchMove(), emcee.moves.DESnookerMove()],
+        vectorize=True,
+    )
+    sampler.run_mcmc(p0, args.nsteps, progress=True, progress_kwargs={"mininterval": 1})
+    samples_raw = sampler.get_chain(flat=True, discard=args.burnin, thin=args.thin)
+    samples = spnta.rescale_samples(samples_raw)
+
+    with open(f"{args.outdir}/samples_raw.npy", "wb") as f:
+        np.save(f, samples_raw)
+    with open(f"{args.outdir}/samples.npy", "wb") as f:
+        np.save(f, samples)

pyvela/pyvela/spnta.py

@@ -165,11 +165,15 @@ def toas(self):
 
     @property
     def param_names(self):
-        return list(vl.get_free_param_names(self.pulsar.model))
+        return np.array(list(vl.get_free_param_names(self.pulsar.model)))
 
     @property
     def param_labels(self):
-        return list(vl.get_free_param_labels(self.pulsar.model))
+        return np.array(list(vl.get_free_param_labels(self.pulsar.model)))
+
+    @property
+    def param_units(self):
+        return np.array(list(vl.get_free_param_units(self.pulsar.model)))
 
     @property
     def scale_factors(self):

pyvela/tests/test_data_files.py

@@ -8,7 +8,7 @@
 
 from pyvela.model import fix_red_noise_components
 from pyvela.parameters import fdjump_rx
-from pyvela.spnta import SPNTA
+from pyvela.spnta import SPNTA, convert_model_and_toas
 from pyvela.vela import vl
 
 # jl.seval("using BenchmarkTools")
@@ -31,7 +31,7 @@
 ]
 
 
-@pytest.fixture(params=datasets, scope="module")
+@pytest.fixture(params=datasets[:2], scope="module")
 def model_and_toas(request):
     dataset = request.param
 
@@ -60,6 +60,15 @@ def model_and_toas(request):
     return spnta, m, t
 
 
+@pytest.mark.parametrize("dataset", datasets[2:])
+def test_read_data(dataset):
+    parfile, timfile = f"{datadir}/{dataset}.par", f"{datadir}/{dataset}.tim"
+    m, t = get_model_and_toas(parfile, timfile, planets=True)
+    model, toas = convert_model_and_toas(m, t)
+    assert len(toas) == len(t)
+    assert len(model.components) <= len(m.components)
+
+
 def test_data(model_and_toas):
     spnta: SPNTA
     spnta, m, t = model_and_toas
@@ -101,6 +110,8 @@ def test_data(model_and_toas):
         and len(spnta.param_labels) == spnta.ndim
     )
 
+    assert len(spnta.param_units) == spnta.ndim
+
 
 def test_chi2(model_and_toas):
     spnta: SPNTA

pyvela/tests/test_run_analysis.py

@@ -5,8 +5,18 @@
 import numpy as np
 
 from pyvela.spnta import SPNTA
+from pyvela import pyvela_script
 from pint.models import get_model_and_toas
 
+prior_str = """
+    {
+        "EFAC": {
+            "distribution": "Uniform",
+            "args": [0.5, 1.5]
+        }
+    }
+"""
+
 
 def test_analysis_NGC6440E_emcee():
     datadir = os.path.dirname(os.path.realpath(__file__)) + "/datafiles"
@@ -15,19 +25,10 @@ def test_analysis_NGC6440E_emcee():
 
     mp, tp = get_model_and_toas(parfile, timfile, planets=True)
 
-    prior_str = StringIO(
-        """{
-            "EFAC": {
-                "distribution": "Uniform",
-                "args": [0.5, 1.5]
-            }
-        }"""
-    )
-
     spnta = SPNTA.from_pint(
         mp,
         tp,
-        custom_priors=prior_str,
+        custom_priors=StringIO(prior_str),
     )
 
     nwalkers = 3 * spnta.ndim
@@ -50,3 +51,27 @@ def test_analysis_NGC6440E_emcee():
         and pint_model_final[pname].uncertainty_value != 0
         for pname in pint_model_final.free_params
     )
+
+
+def test_script_NGC6440():
+    datadir = os.path.dirname(os.path.realpath(__file__)) + "/datafiles"
+    dataset = "NGC6440E"
+    parfile, timfile = f"{datadir}/{dataset}.par", f"{datadir}/{dataset}.tim"
+    outdir = "_NGC6440E_out"
+
+    prior_file = "__prior.json"
+    with open(prior_file, "w") as pf:
+        print(prior_str, file=pf)
+
+    args = f"{parfile} {timfile} -P {prior_file} -o {outdir}".split()
+
+    pyvela_script.main(args)
+
+    assert os.path.isdir(outdir)
+    assert os.path.isfile(f"{outdir}/summary.json")
+    assert os.path.isfile(f"{outdir}/{prior_file}")
+    assert os.path.isfile(f"{outdir}/param_names.txt")
+    assert os.path.isfile(f"{outdir}/param_units.txt")
+    assert os.path.isfile(f"{outdir}/param_scale_factors.txt")
+    assert os.path.isfile(f"{outdir}/samples_raw.npy")
+    assert os.path.isfile(f"{outdir}/samples.npy")

src/model/timing_model.jl

@@ -58,6 +58,7 @@ end
 read_params(model::TimingModel, values) = read_params(model.param_handler, values)
 
 get_free_param_names(model::TimingModel) = get_free_param_names(model.param_handler)
+get_free_param_units(model::TimingModel) = get_free_param_units(model.param_handler)
 get_free_param_labels(model::TimingModel) = get_free_param_labels(model.param_handler)
 
 read_param_values_to_vector(model::TimingModel) =

src/parameter/parameter.jl

@@ -4,6 +4,7 @@ export Parameter,
     ParamHandler,
     read_params,
     get_free_param_names,
+    get_free_param_units,
     get_free_param_labels,
     read_param_values_to_vector,
     get_scale_factors
@@ -173,6 +174,27 @@ function get_free_param_labels(
     return pnames
 end
 
+"""Generate an ordered collection of free parameter units."""
+function get_free_param_units(param_handler::ParamHandler)::Vector{String}
+    pnames = Vector{String}()
+
+    @inbounds for spar in param_handler.single_params
+        if !spar.frozen
+            push!(pnames, spar.original_units)
+        end
+    end
+
+    @inbounds for mpar in param_handler.multi_params
+        for param in mpar.parameters
+            if !param.frozen
+                push!(pnames, param.original_units)
+            end
+        end
+    end
+
+    return pnames
+end
+
 """Generate a collection of free parameter values from a parameter tuple.
 
 Reverse of `read_params()`

test/test_timing_model.jl

@@ -83,5 +83,6 @@
         )
 
         @test get_free_param_names(m2) == ["PHOFF", "F0", "F1", "ECORR1", "ECORR2"]
+        @test length(get_free_param_units(m2)) == length(get_free_param_names(m2))
     end
 end