Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor Configuration and API change #25

Draft
wants to merge 13 commits into
base: master
Choose a base branch
from
Draft

Refactor Configuration and API change #25

wants to merge 13 commits into from

Conversation

eelregit
Copy link
Owner

@eelregit eelregit commented Nov 7, 2023

The units, constants, dtype settings, transfer settings, growth settings, and linear variance settings in Configuration couple more tightly with Cosmology, and should be moved there. After that boltzmann.py can also be independent of Configuration, and functions like jax-cosmo together with cosmology.py.

Likewise, I also try to make Particles more independent of conf. And have some different dtype design for Cosmology and Particles.

To accommodate new features, e.g., sCOLA-like parallelization and observables, the plan is to change the old API

conf = Configuration(...)
cosmo = Cosmology(conf, ...)
modes = white_noise(seed, conf)

def model(modes, ptcl, obsvbl, cosmo, solver):
    cosmo = boltzmann(cosmo, conf)
    modes = linear_modes(modes, cosmo, conf)
    ptcl, obsvbl = lpt(modes, cosmo, conf)
    ptcl, obsvbl = nbody(ptcl, obsvbl, cosmo, conf)
    return obj(obsvbl)

to

def config(...):
    cosmo = Cosmology(...)

    solver = Solver(...)  # the old Configuration, hierarchical, sub-solvers with boundary conditions (BCs)
    modes = white_noise(seed, solver)
    ptcl = pre_init_cond(...)  # or None, to be generated from solver, as additional input to lpt
    obsvbl = Snapshots(...), Lightcone(...)  # observable settings, as additional input to lpt

    return modes, ptcl, obsvbl, cosmo, solver

def model(modes, ptcl, obsvbl, cosmo, solver):
    cosmo = cosmo.prime()  # changed from boltzmann.boltzmann()
    #traverse solver (ptcl and modes too?) pytree (depth/breadth first? flatten is former)
    return obj(traverse(solve(modes, ptcl, obsvbl, cosmo, solver)))

def solve(modes, ptcl, obsvbl, cosmo, solver):
    modes = linear_modes(None, modes, cosmo, solver)  # a=None dependence is moved to front and made explicit
    ptcl, obsvbl, solver = lpt(modes, ptcl, obsvbl, cosmo, solver)  # compute BCs for all sub-solvers
    ptcl, obsvbl, solver = nbody(ptcl, obsvbl, cosmo, solver)  # lpt and nbody also calls obj_imdt -> obsvbl
    obsvbl, solver = obj(obsvbl, solver)  # compute current level obj -> obsvbl
    return obsvbl, solver

Discussions and comments are welcome.

@eelregit eelregit marked this pull request as draft November 7, 2023 01:59
@eelregit eelregit changed the title Refactor Cosmology and Configuration, and API change Refactor Configuration, and API change Nov 7, 2023
@eelregit eelregit changed the title Refactor Configuration, and API change Refactor Configuration and API change Nov 7, 2023
The units, constants, dtype settings, transfer settings, growth
settings, and linear variance settings in Configuration couple more
tightly with Cosmology, and are moved there. After that boltzmann.py can
also be independent of Configuration, and together with cosmology.py
functions like jax-cosmo or other cosmology libraries.
including Cosmology.K, with uppercase K. Omega_k's also changed to
Omega_K's
@eelregit
Copy link
Owner Author

Distances added

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

1 participant