Adding additional prior class

[tsunhopang]

Adding Powerlaw and the Alignedspin prior class; test cases are shown below.

from jax import config
config.update("jax_enable_x64", True)
import jax.numpy as jnp
from jax.random import PRNGKey
from jimgw.prior import Powerlaw
import matplotlib.pyplot as plt

alpha=2
xmin=0.
xmax=100.

key = PRNGKey(42)

prior = Powerlaw(xmin=xmin,
                 xmax=xmax,
                 alpha=alpha,
                 naming=['d_L'])

samples = prior.sample(key, 100000)
log_pdf = prior.log_prob(samples)

samples_sorted = jnp.sort(samples['d_L'])
log_pdf_sorted = log_pdf[samples['d_L'].argsort()]

plt.figure(1)
plt.xlabel('x')
plt.ylabel('pdf')
plt.hist(samples['d_L'], bins=50, density=True)
plt.plot(samples_sorted, jnp.exp(log_pdf_sorted))
plt.yscale('log')
plt.xscale('log')

from jax import config
config.update("jax_enable_x64", True)
import jax.numpy as jnp
from jax.random import PRNGKey
from jimgw.prior import Alignedspin
import matplotlib.pyplot as plt

amax = 0.99
key = PRNGKey(42)

prior = Alignedspin(
    amax=amax,
    naming=['s1_z']
)

samples = prior.sample(key, 100000)
log_pdf = prior.log_prob(samples)

samples_sorted = jnp.sort(samples['s1_z'])
log_pdf_sorted = log_pdf[samples['s1_z'].argsort()]

ref_x = jnp.linspace(-amax, amax, num=10000)
ref_y = -jnp.log(jnp.absolute(ref_x) / amax) / 2 / amax

plt.figure(1)
plt.xlabel('x')
plt.ylabel('pdf')
plt.ylim([0, 3])
plt.hist(samples['s1_z'], bins=50, density=True, label='Samples')
plt.plot(samples_sorted, jnp.exp(log_pdf_sorted), label='Probability')
plt.plot(ref_x, ref_y, label='Reference calculation', linestyle='--')
plt.legend()

[kazewong]

Fixed a couple of minor detail, it is ready to be merged now.

For styling, we can patch it up in a later patch.