Cosmo extension for gwpopulation by HuiTong5 · Pull Request #88 · ColmTalbot/gwpopulation

HuiTong5 · 2024-04-23T03:46:59Z

Try to add cosmology extension for gwpopulation which allows joint cosmology & population analysis, i.e., spectral siren.

The main changes are:

add a cosmology-aware subclass of bilby.hyper.model.Model for cosmology related calculation
have redshift models with changeable cosmology models and corresponding cosmo funcs
make sure the caching step for mass distribution is correct given the situation that the source frame mass samples for spectral siren fit could vary given different cosmological parameters.

…'t use spin

…e over cosmological parameters

ColmTalbot

Hi @HuiTong5 thanks for the contribution! I think this is definitely a change worth having in experimental for a while. I think that until there is a numpy/jax/cupy agnostic cosmology calculator, it should probably stay there.

It looks like the pre-commits are failing, you run these locally following these instructions.

We should also add some tests for this, you can probably base it off the tests for the Redshift models and think about some unit tests for the cosmology conversions.

I'll add that if you would prefer to package this up independently you can see a template here, but that isn't my preference in this case.

gwpopulation/experimental/cosmo_models.py

ColmTalbot · 2024-04-24T12:20:21Z

gwpopulation/experimental/cosmo_models.py

+        if self.astropy_conv == True:
+
+            samples['redshift'] = xp.asarray([z_at_value(cosmo.luminosity_distance, d*u.Mpc,zmax=self.z_max) for d in to_numpy(data['luminosity_distance'])])
+            samples['mass_1'] = data['mass_1']/(1+samples['redshift'])


This conditional block seems overly restrictive and unnecessarily complicated, for example, what if someone is only fitting mass_1? This will raise an error in this case.

I think we should take some time to find a neater solution to this. Two possibilities come to mind:

having a preset list of mass parameters that are checked for, e.g, ["mass_1", "mass_2", "total_mass", "chirp_mass"]. Even better would be to have these be passed as arguments to __init__, something like "redshifted_keys".
for key in data: if key in self.redshift_keys: samples[key] = data[key] / (1 + samples["redshift"]) else: samples[key] = data[key]

looking for a specific flag in keys, i.e., make all redshifted parameters have a _detector suffix, "mass_1_detector" etc. This breaks more strongly from Bilby, but that is already pretty broken by the difference in definition of "mass_1".
for key in data: if key.endswith("_detector"): data[key[:-9]] = data[key] / (1 + data["redshift"])
This can also be combined with the explicit "redshifted_keys" to avoid looping over the entire dictionary and just looking for specific parameters.

Either of these avoid hardcoding any value and in here we just need. Personally, I think the second option is neater as it avoids having the same parameter names meaning different things and potentially also avoides making a new dictionary.

The second option sounds more consistent with your other comments on gwpopulation_pipe. I may try that.

updated in 78835d2 following option 2.

ColmTalbot · 2024-04-24T12:22:43Z

gwpopulation/experimental/cosmo_models.py

+            interp_dl_to_z = splrep(dl,zs,s=0)
+
+            samples['redshift'] = xp.nan_to_num(xp.asarray(splev(to_numpy(data['luminosity_distance']),interp_dl_to_z,ext=0)))
+            samples['mass_1'] = data['mass_1']/(1+samples['redshift'])


Everything from here to L157 is just repeating what was done before, regardless of how we wind up dealing with the implementation of the redshifting, only the following (and equivalent above) should be in the if/else.

zs = to_numpy(self.zs) dl = cosmo.luminosity_distance(to_numpy(zs)).value interp_dl_to_z = splrep(dl,zs,s=0) samples['redshift'] = xp.nan_to_num(xp.asarray(splev(to_numpy(data['luminosity_distance']),interp_dl_to_z,ext=0)))

Fixed in 597fc2c

ColmTalbot · 2024-04-24T12:26:12Z

gwpopulation/experimental/cosmo_models.py

+        samples = dict()
+        if self.astropy_conv == True:
+
+            samples['redshift'] = xp.asarray([z_at_value(cosmo.luminosity_distance, d*u.Mpc,zmax=self.z_max) for d in to_numpy(data['luminosity_distance'])])


z_at_value is now automatically vectorized, so you can remove this list comprehension and just do

Suggested change

samples['redshift'] = xp.asarray([z_at_value(cosmo.luminosity_distance, d*u.Mpc,zmax=self.z_max) for d in to_numpy(data['luminosity_distance'])])

samples['redshift'] = xp.asarray(z_at_value(cosmo.luminosity_distance, to_numpy(data['luminosity_distance']) * u.Mpc, zmax=self.z_max))

We can also possibly avoid repeated transfer between array types by setting

data["luminosity_distance"] = to_numpy(data["luminosity_distance"]) somewhere.

When fitting cosmology I think this will cause a slowdown when transferring computations between astropy and gwpopulation. I think this class should include a function like (assuming flatLambdaCDM):

def comoving_distance(z): E_z = xp.sqrt(Om * (1+self.zs)**3 + (1-Om)) return c/H0 * cumtrapz(Ez, 0, z) def luminosity_distance(z): return (1+z) * self.comoving_distance(z)

This way these transformations can be handled all natively in the CPU/GPU agnostic manner. To be generic there could be an option to have these computed with astropy.

z_at_value comment fixed in 597fc2c.

Indeed, not only this conversion, but also some other calculations by astropy may introduce a significant slowdown. I've implemented some interpolation to replace astropy.cosmo call somewhere else. I will have a more comprehensive thinking of possible approximation of all the astropy calculations in this PR.

ColmTalbot · 2024-04-24T12:29:14Z

gwpopulation/experimental/cosmo_models.py

+
+        return samples
+
+    def dL_by_dz(self, z, dl, **parameters):


dl should be included in the docstring.

Also, please use full variable names, redshift, luminosity_distance for readability. The method name is probably fine, but luminosity distance_to_redshift_jacobian would be better.

I've just noticed that this is taking both the redshift and distance when there should be a deterministic mapping between them. I worry that this is inviting user error.

What if you move the logic to convert from distance to redshift to inside this method (or another method that calls this) and then return both redshift and dl_by_dz (lowercase for variable names)?

docstring and variables names fixed in 597fc2c.

updated in 78835d where I combine the conversion of samples from detector frame to source frame and the calculation of Jabobian term in one func which only need detector frame samples as input.

ColmTalbot · 2024-04-24T12:31:19Z

gwpopulation/experimental/cosmo_models.py

+    def dL_by_dz(self, z, dl, **parameters):
+
+        """
+        Calculates the detector frame to source frame Jacobian d_det/d_sour for dL and z


Again, we don't need to use abbreviations here, a longer docstring will be helpful. Ideally, with an equation, there are docs here that explain some of the subtleties of getting this to render correctly in the online docs.

gwpopulation/experimental/cosmo_models.py

gwpopulation/models/mass.py

gwpopulation/experimental/cosmo_models.py

Co-authored-by: Colm Talbot <talbotcolm@gmail.com>

codecov-commenter · 2024-04-30T08:22:30Z

Codecov Report

Attention: Patch coverage is 66.66667% with 4 lines in your changes are missing coverage. Please review.

Project coverage is 94.18%. Comparing base (b7edefa) to head (5ab859e).

Files	Patch %	Lines
gwpopulation/models/mass.py	66.67%	2 Missing ⚠️
gwpopulation/vt.py	66.67%	2 Missing ⚠️

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #88      +/-   ##
==========================================
- Coverage   94.74%   94.18%   -0.56%     
==========================================
  Files          11       11              
  Lines         799      808       +9     
==========================================
+ Hits          757      761       +4     
- Misses         42       47       +5

Flag	Coverage Δ
python3.10	`94.06% <66.67%> (-0.56%)`	⬇️
python3.11	`94.06% <66.67%> (-0.56%)`	⬇️
python3.9	`94.05% <66.67%> (-0.56%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

ColmTalbot · 2024-05-07T15:57:33Z

@afarah18 has some code based on this paper that we might be able to adapt to handle the cosmological conversions. It isn't as flexible as astropy, it currently just handles variable H0 and Omega_m, but I suspect it can be extended to w quite easily.

…hift class

HuiTong5 · 2024-05-11T08:37:01Z

Thanks for sharing! I will take a look at the paper and Amenda's code when I have time. I have pushed some updates based on our previous discussion. I haven't done any actual calculations/testing, so there may be a lot of bugs in it. I'll do some testing soon. But I think the overall framework has improved to some extent.

ColmTalbot · 2024-05-11T18:26:43Z

The new changes generally look good! Amanda and I are working on packaging up the other cosmology implementation and I'll probably match the astropy API closely so it should be a drop in replacement that I can do later.

gwpopulation/experimental/cosmo_models.py

gwpopulation/vt.py

…anyone would like to turn it on given the enormous time cost. Can add astropy_conv back later with other cosmology implementation. Aslo, minor fix for vt surveyed_hypervolume

ColmTalbot · 2024-05-14T02:13:10Z

I made some changes to the way the modules that need the backend to change are specified, the short version is that you should now add the new module to the setup.cfg file instead of the backend.py.

I also pushed a version of the cosmology implementation that @afarah18 and I put together in https://github.com/ColmTalbot/wcosmo. It will be a fairly simple drop-in change once this has been merged and I added some similar code to what you have in this PR so that I can use it as a free-standing version in the meantime.

HuiTong5 · 2024-05-14T04:47:06Z

The new repo for cosmology implementation looks great! I noticed there content in models.py in the new repo is almost the same as the new module cosmo_models.py in this PR. What's the plan for this PR? Should I call the functions in the new repo to avoid repetitive code?

ColmTalbot · 2024-05-14T12:46:25Z

The new repo for cosmology implementation looks great! I noticed there content in models.py in the new repo is almost the same as the new module cosmo_models.py in this PR. What's the plan for this PR? Should I call the functions in the new repo to avoid repetitive code?

That repo is essentially a holding place so that I can use the code while we get the changes into gwpopulation. I should have added an attribution to you in there as it was strongly influenced by this PR, sorry! I'll do that today. My plan is:

finalize and merge this PR using astropy as currently
I open a new PR to use the cosmology code in the new repo either natively or as an external import

The latter stage would involve removing models.py basically entirely.

HuiTong5 · 2024-05-15T08:49:51Z

I had a quick test on the time cost for single likelihood evaluation on the head node of pcdev3. The cosmo one costed roughly 0.53s and non cosmo one costed roughly 0.15s.

ColmTalbot · 2024-05-15T12:17:04Z

I had a quick test on the time cost for single likelihood evaluation on the head node of pcdev3. The cosmo one costed roughly 0.53s and non cosmo one costed roughly 0.15s.

Nice! Which backend is that using?

~~I'm going to merge this now, so I can make an attempt at implementing the other cosmology code. Thanks @HuiTong5!~~ EDIT: I just saw this is marked as draft, so I'll hold off until you remove that label.

HuiTong5 · 2024-05-15T12:20:19Z

I was using cupy

HuiTong5 · 2024-05-15T12:25:28Z

~~I'm going to merge this now, so I can make an attempt at implementing the other cosmology code. Thanks @HuiTong5!~~ EDIT: I just saw this is marked as draft, so I'll hold off until you remove that label.

I don't think I plan to have any other changes if you are happy with this. What label should I have? Should I have DEV or I just need to simply remove the "Draft".

ColmTalbot · 2024-05-15T12:35:09Z

Great, I removed Draft from the issue title, in future, there is a "mark as draft" button on PRs somewhere that makes it more clear.

Hui Tong and others added 13 commits April 1, 2024 22:07

add cosmo likelihood, cosmo redshift model and cosmo vt func

3b6cbc6

update backend for new cosmo funcs

b9ae3ef

minor fix

4b86a3f

fix sample conversion from detector frame to source frame when we don…

f1a54aa

…'t use spin

correct caching process for cosmo fit when dataset[mass_1] will chang…

e1a1a21

…e over cosmological parameters

minor fix

e28f6d3

fix for jacobian term dL_by_dz

7c6bb07

replace redshift model astropy dvc_dz calculation by interpolation

d4e79f1

minor fix

92cce59

fix for jacobian term and clean redshift model

f3d55ab

FLatwCDM model

735cdb5

flexible cosmo model input

0f0ed97

replace cosmolikelihood and cosmo vt with cosmo model to simplify things

bd3c53c

ColmTalbot requested changes Apr 24, 2024

View reviewed changes

jacobgolomb reviewed Apr 24, 2024

View reviewed changes

gwpopulation/experimental/cosmo_models.py Outdated Show resolved Hide resolved

Hui Tong and others added 2 commits April 30, 2024 00:12

pre commit fix

e4a0562

remove "inheriting from object"

6f0043f

Co-authored-by: Colm Talbot <talbotcolm@gmail.com>

Hui Tong and others added 2 commits April 30, 2024 21:00

fixes for code format. Avoid some reuse and duplication

597fc2c

minor fix

534771c

Hui Tong added 5 commits May 11, 2024 00:31

create CosmoMixin for cosmo related calculation to clean up CosmoReds…

78835d2

…hift class

add cache flag for BaseSmoothedMassDistribution

228b410

minor fix for mass cache

7b030f5

create TestRedshift as potential base redshift model class model

62c6f0f

minor fix for varibale names in TestRedshift class

832d614

ColmTalbot reviewed May 11, 2024

View reviewed changes

reorganize redshift model class

7822324

minor fix

dcf2b1e

ColmTalbot reviewed May 12, 2024

View reviewed changes

gwpopulation/vt.py Outdated Show resolved Hide resolved

Hui Tong added 4 commits May 12, 2024 21:16

minor fix

e4eee04

remove astropy_conv flag temporarily because it's very unlikely that …

1ead1e2

…anyone would like to turn it on given the enormous time cost. Can add astropy_conv back later with other cosmology implementation. Aslo, minor fix for vt surveyed_hypervolume

minor fix

5a45f38

minor fix

0ec4778

Hui Tong added 2 commits May 13, 2024 21:30

Merge branch 'main' into cosmo_dev

33cad09

update the backend specification for cosmo_model in setup.cfg

5ab859e

ColmTalbot changed the title ~~Draft: Cosmo extension for gwpopulation~~ Cosmo extension for gwpopulation May 15, 2024

ColmTalbot merged commit 261f2b8 into ColmTalbot:main May 15, 2024

ColmTalbot mentioned this pull request May 15, 2024

backend agnostic cosmology #92

Merged

	samples['redshift'] = xp.asarray([z_at_value(cosmo.luminosity_distance, d*u.Mpc,zmax=self.z_max) for d in to_numpy(data['luminosity_distance'])])
	samples['redshift'] = xp.asarray(z_at_value(cosmo.luminosity_distance, to_numpy(data['luminosity_distance']) * u.Mpc, zmax=self.z_max))

Conversation

HuiTong5 commented Apr 23, 2024

Uh oh!

ColmTalbot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

codecov-commenter commented Apr 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

ColmTalbot commented May 7, 2024

Uh oh!

HuiTong5 commented May 11, 2024

Uh oh!

ColmTalbot commented May 11, 2024

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ColmTalbot commented May 14, 2024

Uh oh!

HuiTong5 commented May 14, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ColmTalbot commented May 14, 2024

Uh oh!

HuiTong5 commented May 15, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ColmTalbot commented May 15, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HuiTong5 commented May 15, 2024

Uh oh!

HuiTong5 commented May 15, 2024

Uh oh!

ColmTalbot commented May 15, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

codecov-commenter commented Apr 30, 2024 •

edited

Loading

HuiTong5 commented May 14, 2024 •

edited

Loading

HuiTong5 commented May 15, 2024 •

edited

Loading

ColmTalbot commented May 15, 2024 •

edited

Loading