Frequently Asked Questions¶

We are collecting answers to frequently asked questions here. Contributions welcome! Besides the questions below, you can also check the following resources:

How to write my own population model?¶

GWKokab is designed to be flexible and user-friendly. To create your own population model, you can subclass the Distribution class. If you need to generate a population and run inference, you should implement both the sample and log_prob methods. However, if you only need to run inference, you can just implement the log_prob method.

Here is an example of a simple population model:

>>> import numpyro
>>> from jax import random as jrd
>>> from numpyro.distributions import Distribution
>>>
>>> class MyPopulationModel(Distribution):
...     def __init__(self, loc, scale):
...         self.loc = loc
...         self.scale = scale
...
...     def sample(self, key, sample_shape=()):
...         return numpyro.distributions.Normal(self.loc, self.scale).sample(
...             key, sample_shape
...         )
...
...     def log_prob(self, value):
...         return numpyro.distributions.Normal(self.loc, self.scale).log_prob(value)
>>>
>>> samples = MyPopulationModel(0.0, 1.0).sample(jrd.key(0), (1000,))
>>> samples.shape
(1000,)

GWKokab is slow on GPU¶

There can be several reasons why GWKokab is slow on GPU. Here are some common issues:

JAX is not installed correctly. Make sure you have the correct version of JAX installed along with the necessary dependencies. Check installation instructions in the GWKokab’s documentation and the JAX’s documentation
Use appropriate environment variables. Following are the common environment variables GWKokab dev team uses,
```
export NPROC=16
export intra_op_parallelism_threads=4
export OPENBLAS_NUM_THREADS=4
export TF_CPP_MIN_LOG_LEVEL=1
export XLA_PYTHON_CLIENT_ALLOCATOR=platform
export TF_FORCE_GPU_ALLOW_GROWTH=false
export JAX_COMPILATION_CACHE_DIR="/tmp/jax_cache"
```
Their values are adjusted based on the system configuration and hardware. You may need to adjust them based on your system configuration. Therefore it is recommended to experiment with different values to find the best configuration for your system and check their documentation for more information.
See following articles,