Skip to content

powerlaw

broken_power(rs, condlist, rbins, amps, pows, c)

Function which returns a broken powerlaw evaluated at rs.

Parameters:

rs : jax.Array Array of rs at which to compute pl. condlist : tuple tuple which enocdes which rs are evaluated by which parametric function rbins : jax.Array Array of bin edges for power laws amps : jax.Array Amplitudes of power laws pows : jax.Array Exponents of power laws c : float Constant offset for powerlaws

Source code in witch/powerlaw.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
@jax.jit
def broken_power(
    rs: jax.Array,
    condlist: tuple,
    rbins: jax.Array,
    amps: jax.Array,
    pows: jax.Array,
    c: float,
) -> jax.Array:
    """
    Function which returns a broken powerlaw evaluated at rs.

    Parameters:
    -----------
    rs : jax.Array
        Array of rs at which to compute pl.
    condlist : tuple
        tuple which enocdes which rs are evaluated by which parametric function
    rbins : jax.Array
        Array of bin edges for power laws
    amps : jax.Array
        Amplitudes of power laws
    pows : jax.Array                                                                                                                                                                                                                                                                            Exponents of power laws
    c : float
        Constant offset for powerlaws
    """
    cur_c = c  # TODO: necessary?
    funclist = []
    for i in range(len(condlist) - 1, -1, -1):
        funclist.append(
            partial(power, rbin=rbins[i + 1], cur_amp=amps[i], cur_pow=pows[i], c=cur_c)
        )
        cur_c += amps[i] * (rbins[i] ** pows[i] - rbins[i + 1] ** pows[i])
    return jnp.piecewise(rs, condlist, funclist)

power(x, rbin, cur_amp, cur_pow, c)

Function which returns the powerlaw, given the bin-edge constraints. Exists to be partialed.

Parameters:

x : float Dummy variable to be partialed over rbin : float Edge of bin for powerlaw cur_amp : float Amplitude of power law cur_pow : float Power of power law c : float Constant offset

Returns:

Name Type Description
tmp float

Powerlaw evaluated at x
Source code in witch/powerlaw.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
@jax.jit
def power(x: float, rbin: float, cur_amp: float, cur_pow: float, c: float):
    """
    Function which returns the powerlaw, given the bin-edge constraints. Exists to be partialed.

    Parameters:
    -----------
    x : float
        Dummy variable to be partialed over
    rbin : float
        Edge of bin for powerlaw
    cur_amp : float
        Amplitude of power law
    cur_pow : float
        Power of power law
    c : float
        Constant offset

    Returns
    -------
    tmp : float
        Powerlaw evaluated at x
    """
    tmp = cur_amp * (x**cur_pow - rbin**cur_pow) + c
    return tmp

profile_to_broken_power(rs, ys, condlist, rbins)

Estimates a non-parametric broken power profile from a generic profile. Note this is an estimation only; in partciular since we fit piece-wise the c's get messed up. This broken powerlaw should then be fit to the data.

Parameters:

Name Type Description Default
rs ArrayLike

Array of radius values for the profile

 required
ys ArrayLike

Profile y values

 required
condlist list[ArrayLike]

List which defines which powerlaws map to which radii. See broken_power

 required
rbins ArrayLike

Array of bin edges defining the broken powerlaws

 required

Returns:

Name Type Description
amps array

Best fit amps for the powerlaws
pows array

Best fit powers for the powerlaws
c float

Best fit c for only the outermost powerlaw
Source code in witch/powerlaw.py 
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
def profile_to_broken_power(
    rs: ArrayLike, ys: ArrayLike, condlist: list[jax.Array], rbins: ArrayLike
) -> tuple[jax.Array, jax.Array, float]:
    """
    Estimates a non-parametric broken power profile from a generic profile.
    Note this is an estimation only; in partciular since we fit piece-wise
    the c's get messed up. This broken powerlaw should then be fit to the
    data.

    Parameters
    ----------
    rs : ArrayLike
        Array of radius values for the profile
    ys : ArrayLike
        Profile y values
    condlist : list[ArrayLike]
        List which defines which powerlaws map to which radii. See broken_power
    rbins : ArrayLike
        Array of bin edges defining the broken powerlaws

    Returns
    -------
    amps : jnp.array
        Best fit amps for the powerlaws
    pows : jnp.array
        Best fit powers for the powerlaws
    c : float
        Best fit c for only the outermost powerlaw
    """
    rs = jnp.array([x if x != 0 else 1e-1 for x in rs])  # Dont blow up

    rbins = jnp.array(
        [x if x != 0 else jnp.amin(rs) for x in rbins]
    )  # Dont blow up 2.0

    amps = jnp.zeros(len(condlist))
    pows = jnp.zeros(len(condlist))

    for i in range(len(condlist)):
        xdata = rs[condlist[i]]
        ydata = ys[condlist[i]]
        if i == len(condlist) - 1:
            popt, pcov = curve_fit(power, xdata, ydata, method="trf")
        else:
            popt, pcov = curve_fit(
                power,
                xdata,
                ydata,
                method="trf",
                p0=[rbins[::-1][i], np.amax(ydata) * 1e5, -2, 0.0],
            )
        if i == 0:
            c = popt[3]
        amps = amps.at[i].set(popt[1])
        pows = pows.at[i].set(popt[2])

    return amps[::-1], pows[::-1], c