diff --git a/.gitignore b/.gitignore index 025b5cd3..bc7f8898 100644 --- a/.gitignore +++ b/.gitignore @@ -188,4 +188,8 @@ cython_debug/ profile.pb.gz -*_results.pdf \ No newline at end of file +*_results.pdf + +pyvela/tests/_*_out/ +pyvela_results +__prior.json \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index ce73a53d..c8dc2965 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 4d1de122..19ff16f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,6 @@ maintainers = [ {name = "Abhimanyu Susobhanan", email = "abhisrkckl@gmail.com"} ] - 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"} diff --git a/pyvela/pyvela/pyvela_script.py b/pyvela/pyvela/pyvela_script.py new file mode 100644 index 00000000..fc401bb4 --- /dev/null +++ b/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) diff --git a/pyvela/pyvela/spnta.py b/pyvela/pyvela/spnta.py index 6dd74045..d8879fe2 100644 --- a/pyvela/pyvela/spnta.py +++ b/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): diff --git a/pyvela/tests/test_data_files.py b/pyvela/tests/test_data_files.py index 0b4f4b1a..d15edaf9 100644 --- a/pyvela/tests/test_data_files.py +++ b/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 diff --git a/pyvela/tests/test_run_analysis.py b/pyvela/tests/test_run_analysis.py index ab978168..36620733 100644 --- a/pyvela/tests/test_run_analysis.py +++ b/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") diff --git a/src/model/timing_model.jl b/src/model/timing_model.jl index 969025a2..2bc3e5b7 100644 --- a/src/model/timing_model.jl +++ b/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) = diff --git a/src/parameter/parameter.jl b/src/parameter/parameter.jl index 4284bb6c..111b6fde 100644 --- a/src/parameter/parameter.jl +++ b/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()` diff --git a/test/test_timing_model.jl b/test/test_timing_model.jl index 54eed216..e5645e81 100644 --- a/test/test_timing_model.jl +++ b/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