def measure_cd(x, sr, start=-0.25, end=0.25, bins=2000, wavlen=1550e-9):
'''
References:
Zhou, H., Li, B., Tang, et. al, 2016. Fractional fourier transformation-based
blind chromatic dispersion estimation for coherent optical communications.
Journal of Lightwave Technology, 34(10), pp.2371-2380.
'''
c = 299792458.
p = jnp.linspace(start, end, bins)
N = x.shape[0]
K = p.shape[0]
L = jnp.zeros(K, dtype=jnp.float32)
def f(_, pi):
return None, jnp.sum(
jnp.abs(xop.frft(jnp.abs(xop.frft(x, pi))**2, -1))**2)
# Use `scan` instead of `vmap` here to avoid potential large memory allocation.
# Despite the speed of `scan` scales surprisingly well to large bins,
# the speed has a lowerbound e.g 600ms at bins=1, possiblely related to the blind
# migration of `frft` from Github :) (could frft be jitted in theory?).
# TODO review `frft`
_, L = xop.scan(f, None, p)
B2z = jnp.tan(jnp.pi / 2 -
(p - 1) / 2 * jnp.pi) / (sr * 2 * jnp.pi / N * sr)
Dz_set = -B2z / wavlen**2 * 2 * jnp.pi * c # the swept set of CD metrics
Dz_hat = Dz_set[jnp.argmin(L)] # estimated accumulated CD
return Dz_hat, L, Dz_set
def _rvs(self, a, b):
# We use inverse transform method:
# z ~ ppf(U), where U ~ Uniform(cdf(a), cdf(b)).
# ~ Uniform(arctan(a), arctan(b)) / pi + 1/2
u = random.uniform(self._random_state, shape=self._size,
minval=np.arctan(a), maxval=np.arctan(b))
return np.tan(u)
def f(c, a): assert a.shape == (3,) assert c.shape == (4,) b = np.cos(np.sum(np.sin(a)) + np.sum(np.cos(c)) + np.sum(np.tan(d))) c = np.sin(c * b) assert b.shape == () return c, b
def sample(self, key, sample_shape=()):
# We use inverse transform method:
# z ~ inv_cdf(U), where U ~ Uniform(cdf(low), cdf(high)).
# ~ Uniform(arctan(low), arctan(high)) / pi + 1/2
size = sample_shape + self.batch_shape
minval = -np.arctan(self.base_loc)
maxval = np.pi / 2
u = minval + random.uniform(key, shape=size) * (maxval - minval)
return self.base_loc + np.tan(u)
def tick_buffer(carry, X):
params, state = carry
K = jnp.tan(jnp.pi * params["cutoff"])
norm = 1.0 / (1.0 + K / params["resonance"] + K * K)
state["a0"] = K * K * norm
state["a1"] = 2.0 * state["a0"]
state["a2"] = state["a0"]
state["b1"] = 2.0 * (K * K - 1) * norm
state["b2"] = (1.0 - K / params["resonance"] + K * K) * norm
return lax.scan(tick, carry, X)
def inverse_bearings(observation, s1, s2):
"""
Inverse the bearings observation to the location as if there was no noise,
This is only used to provide an initial point for the linearization of the IEKS and ICKS.
Parameters
----------
observation: (2) array
The bearings observation
s1: (2) array
The first sensor position
s2: (2) array
The second sensor position
Returns
-------
out: (2) array
The inversed position of the state
"""
tan_theta = jnp.tan(observation)
A = jnp.array([[tan_theta[0], -1], [tan_theta[1], -1]])
b = jnp.array([s1[0] * tan_theta[0] - s1[1], s2[0] * tan_theta[1] - s2[1]])
return jnp.linalg.solve(A, b)
def _ppf(self, q):
return jnp.tan(jnp.pi * q - jnp.pi / 2.0)
def _ppf(self, q):
return jnp.tan(jnp.pi / 2 * q)
# Create TFC class:
myTfc = utfc(n, nC, m, x0=th0, xf=thf)
th = myTfc.x
H = myTfc.H
# Create constrained expression:
g = lambda th, xi: np.dot(H(th), xi)
r = lambda th,xi: g(th,xi)+\
(th-thf)/(th0-thf)*(r0-g(th0*np.ones_like(th),xi))+\
(th-th0)/(thf-th0)*(rf-g(thf*np.ones_like(th),xi))
# Create loss function:
dr = egrad(r)
d2r = egrad(dr)
L = lambda xi: -r(th,xi)**2*(dr(th,xi)*np.tan(th)+2.*d2r(th,xi))+\
-np.tan(th)*dr(th,xi)**3+3.*r(th,xi)*dr(th,xi)**2+r(th,xi)**3
# Solve the problem:
xi = np.zeros(H(th).shape[1])
xi, _, time = NLLS(xi, L, timer=True)
# Print out statistics:
print("Solution time: {0} seconds".format(time))
# Plot the solution and residual
p = MakePlot([r"$y$"], [r"$x$"])
p.ax[0].plot(r(th, xi) * np.sin(th), r(th, xi) * np.cos(th), "k")
p.ax[0].axis("equal")
p.ax[0].grid(True)
p.ax[0].invert_yaxis()
def tan(a: Numeric):
return jnp.tan(a)
def _rvs(self):
# TODO: move this implementation upstream to jax.random.standard_cauchy
# Another way is to generate X, Y ~ Normal(0, 1) and return X / Y
u = random.uniform(self._random_state, shape=self._size)
return np.tan(np.pi * (u - 0.5))
def _ppf(self, q):
return np.tan(np.pi / 2 * q)
def tan(x): if isinstance(x, JaxArray): x = x.value return JaxArray(jnp.tan(x))
def reflected_phase_curve_inhomogeneous(phases, omega_0, omega_prime, x1, x2,
A_g, a_rp):
"""
Reflected light phase curve for an inhomogeneous sphere by
Heng, Morris & Kitzmann (2021), with inspiration from Hu et al. (2015).
Parameters
----------
phases : `~np.ndarray`
Orbital phases of each observation defined on (0, 1)
omega_0 : tensor-like
Single-scattering albedo of the less reflective region.
Defined on (0, 1).
omega_prime : tensor-like
Additional single-scattering albedo of the more reflective region,
such that the single-scattering albedo of the reflective region is
``omega_0 + omega_prime``. Defined on (0, ``1-omega_0``).
x1 : tensor-like
Start longitude of the darker region [radians] on (-pi/2, pi/2)
x2 : tensor-like
Stop longitude of the darker region [radians] on (-pi/2, pi/2)
a_rp : float, tensor-like
Semimajor axis scaled by the planetary radius
Returns
-------
flux_ratio_ppm : tensor-like
Flux ratio between the reflected planetary flux and the stellar flux
in units of ppm.
g : tensor-like
Scattering asymmetry factor on (-1, 1)
q : tensor-like
Integral phase function
"""
g = _g_from_ag(A_g, omega_0, omega_prime, x1, x2)
# Redefine alpha to be on (-pi, pi)
alpha = (2 * np.pi * phases - np.pi).astype(floatX)
abs_alpha = np.abs(alpha).astype(floatX)
alpha_sort_order = np.argsort(alpha)
sin_abs_sort_alpha = np.sin(abs_alpha[alpha_sort_order]).astype(floatX)
sort_alpha = alpha[alpha_sort_order].astype(floatX)
# Equation 34 for Henyey-Greestein
P_star = (1 - g**2) / (1 + g**2 + 2 * g * jnp.cos(abs_alpha))**1.5
# Equation 36
P_0 = (1 - g) / (1 + g)**2
Rho_S, Rho_S_0, Rho_L, Rho_C = rho(omega_0, P_0, P_star)
Rho_S_prime, Rho_S_0_prime, Rho_L_prime, Rho_C_prime = rho(
omega_prime, P_0, P_star)
alpha_plus = jnp.sin(abs_alpha / 2) + jnp.cos(abs_alpha / 2)
alpha_minus = jnp.sin(abs_alpha / 2) - jnp.cos(abs_alpha / 2)
# Equation 11:
Psi_0 = jnp.where((alpha_plus > -1) & (alpha_minus < 1),
jnp.log((1 + alpha_minus) * (alpha_plus - 1) /
(1 + alpha_plus) / (1 - alpha_minus)), 0)
Psi_S = 1 - 0.5 * (jnp.cos(abs_alpha / 2) -
1.0 / jnp.cos(abs_alpha / 2)) * Psi_0
Psi_L = (jnp.sin(abs_alpha) +
(np.pi - abs_alpha) * jnp.cos(abs_alpha)) / np.pi
Psi_C = (-1 + 5 / 3 * jnp.cos(abs_alpha / 2)**2 -
0.5 * jnp.tan(abs_alpha / 2) * jnp.sin(abs_alpha / 2)**3 * Psi_0)
# Table 1:
condition_a = (-np.pi / 2 <= alpha - np.pi / 2)
condition_0 = ((alpha - np.pi / 2 <= np.pi / 2) & (np.pi / 2 <= alpha + x1)
& (alpha + x1 <= alpha + x2))
condition_1 = ((alpha - np.pi / 2 <= alpha + x1) &
(alpha + x1 <= np.pi / 2) & (np.pi / 2 <= alpha + x2))
condition_2 = ((alpha - np.pi / 2 <= alpha + x1) &
(alpha + x1 <= alpha + x2) & (alpha + x2 <= np.pi / 2))
condition_b = (alpha + np.pi / 2 <= np.pi / 2)
condition_3 = ((alpha + x1 <= alpha + x2) & (alpha + x2 <= -np.pi / 2) &
(-np.pi / 2 <= alpha + np.pi / 2))
condition_4 = ((alpha + x1 <= -np.pi / 2) & (-np.pi / 2 <= alpha + x2) &
(alpha + x2 <= alpha + np.pi / 2))
condition_5 = ((-np.pi / 2 <= alpha + x1) & (alpha + x1 <= alpha + x2) &
(alpha + x2 <= alpha + np.pi / 2))
integration_angles = [
[alpha - np.pi / 2, np.pi / 2], [alpha - np.pi / 2, alpha + x1],
[alpha - np.pi / 2, alpha + x1, alpha + x2, np.pi / 2],
[-np.pi / 2, alpha + np.pi / 2], [alpha + x2, alpha + np.pi / 2],
[-np.pi / 2, alpha + x1, alpha + x2, alpha + np.pi / 2]
]
conditions = [
condition_a & condition_0,
condition_a & condition_1,
condition_a & condition_2,
condition_b & condition_3,
condition_b & condition_4,
condition_b & condition_5,
]
Psi_S_prime = 0
Psi_L_prime = 0
Psi_C_prime = 0
for condition_i, angle_i in zip(conditions, integration_angles):
for i, phi_i in enumerate(angle_i):
sign = (-1)**(i + 1)
I_phi_S, I_phi_L, I_phi_C = I(alpha, phi_i)
Psi_S_prime += jnp.where(condition_i, sign * I_phi_S, 0)
Psi_L_prime += jnp.where(condition_i, sign * I_phi_L, 0)
Psi_C_prime += jnp.where(condition_i, sign * I_phi_C, 0)
# Compute everything for alpha=0
angles_alpha0 = [-np.pi / 2, x1, x2, np.pi / 2]
Psi_S_prime_alpha0 = 0
Psi_L_prime_alpha0 = 0
Psi_C_prime_alpha0 = 0
for i, phi_i in enumerate(angles_alpha0):
sign = (-1)**(i + 1)
I_phi_S_alpha0, I_phi_L_alpha0, I_phi_C_alpha0 = I(0, phi_i)
Psi_S_prime_alpha0 += sign * I_phi_S_alpha0
Psi_L_prime_alpha0 += sign * I_phi_L_alpha0
Psi_C_prime_alpha0 += sign * I_phi_C_alpha0
# Equation 37
F_S = np.pi / 16 * (omega_0 * Rho_S * Psi_S +
omega_prime * Rho_S_prime * Psi_S_prime)
F_L = np.pi / 12 * (omega_0 * Rho_L * Psi_L +
omega_prime * Rho_L_prime * Psi_L_prime)
F_C = 3 * np.pi / 64 * (omega_0 * Rho_C * Psi_C +
omega_prime * Rho_C_prime * Psi_C_prime)
sobolev_fluxes = F_S + F_L + F_C
F_max = jnp.max(sobolev_fluxes)
Psi = sobolev_fluxes / F_max
flux_ratio_ppm = 1e6 * a_rp**-2 * Psi * A_g
q = _integral_phase_function(Psi, sin_abs_sort_alpha, sort_alpha,
alpha_sort_order)
# F_0 = F_S_alpha0 + F_L_alpha0 + F_C_alpha0
return flux_ratio_ppm, g, q
def reflected_phase_curve(phases, omega, g, a_rp):
"""
Reflected light phase curve for a homogeneous sphere by
Heng, Morris & Kitzmann (2021).
Parameters
----------
phases : `~np.ndarray`
Orbital phases of each observation defined on (0, 1)
omega : tensor-like
Single-scattering albedo as defined in
g : tensor-like
Scattering asymmetry factor, ranges from (-1, 1).
a_rp : float, tensor-like
Semimajor axis scaled by the planetary radius
Returns
-------
flux_ratio_ppm : tensor-like
Flux ratio between the reflected planetary flux and the stellar flux in
units of ppm.
A_g : tensor-like
Geometric albedo derived for the planet given {omega, g}.
q : tensor-like
Integral phase function
"""
# Convert orbital phase on (0, 1) to "alpha" on (0, np.pi)
alpha = jnp.asarray(2 * np.pi * phases - np.pi)
abs_alpha = jnp.abs(alpha)
alpha_sort_order = jnp.argsort(alpha)
sin_abs_sort_alpha = jnp.sin(abs_alpha[alpha_sort_order])
sort_alpha = alpha[alpha_sort_order]
gamma = jnp.sqrt(1 - omega)
eps = (1 - gamma) / (1 + gamma)
# Equation 34 for Henyey-Greestein
P_star = (1 - g**2) / (1 + g**2 + 2 * g * jnp.cos(alpha))**1.5
# Equation 36
P_0 = (1 - g) / (1 + g)**2
# Equation 10:
Rho_S = P_star - 1 + 0.25 * ((1 + eps) * (2 - eps))**2
Rho_S_0 = P_0 - 1 + 0.25 * ((1 + eps) * (2 - eps))**2
Rho_L = 0.5 * eps * (2 - eps) * (1 + eps)**2
Rho_C = eps**2 * (1 + eps)**2
alpha_plus = jnp.sin(abs_alpha / 2) + jnp.cos(abs_alpha / 2)
alpha_minus = jnp.sin(abs_alpha / 2) - jnp.cos(abs_alpha / 2)
# Equation 11:
Psi_0 = jnp.where((alpha_plus > -1) & (alpha_minus < 1),
jnp.log((1 + alpha_minus) * (alpha_plus - 1) /
(1 + alpha_plus) / (1 - alpha_minus)), 0)
Psi_S = 1 - 0.5 * (jnp.cos(abs_alpha / 2) -
1.0 / jnp.cos(abs_alpha / 2)) * Psi_0
Psi_L = (jnp.sin(abs_alpha) +
(np.pi - abs_alpha) * jnp.cos(abs_alpha)) / np.pi
Psi_C = (-1 + 5 / 3 * jnp.cos(abs_alpha / 2)**2 -
0.5 * jnp.tan(abs_alpha / 2) * jnp.sin(abs_alpha / 2)**3 * Psi_0)
# Equation 8:
A_g = omega / 8 * (P_0 - 1) + eps / 2 + eps**2 / 6 + eps**3 / 24
# Equation 9:
Psi = ((12 * Rho_S * Psi_S + 16 * Rho_L * Psi_L + 9 * Rho_C * Psi_C) /
(12 * Rho_S_0 + 16 * Rho_L + 6 * Rho_C))
flux_ratio_ppm = 1e6 * (a_rp**-2 * A_g * Psi)
q = _integral_phase_function(Psi, sin_abs_sort_alpha, sort_alpha,
alpha_sort_order)
return flux_ratio_ppm, A_g, q
def _tan(x):
return jnp.tan(x)
def _logsnr_schedule_cosine(t, *, logsnr_min, logsnr_max):
b = onp.arctan(onp.exp(-0.5 * logsnr_max))
a = onp.arctan(onp.exp(-0.5 * logsnr_min)) - b
return -2. * jnp.log(jnp.tan(a * t + b))
def state_dynamics(state, control, *params):
# sanity
assert len(state.shape) == len(control.shape) == 1
assert state.shape[-1] == Fossen.state_dim
assert control.shape[-1] == Fossen.control_dim
# constant parameters
Nrr, Izz, Kt1, zb, Mqq, ycp, xb, zcp, Yvv, yg, Ixx, Kt0, Xuu, xg, Zww, W, m, B, zg, Kpp, Qt1, Qt0, Iyy, yb, xcp = params
# extract state and control
# phi, theta, psi = roll, pitch, yaw
x, y, z, phi, theta, psi, u, v, w, p, q, r = state
rpm0, rpm1, de, dr = control
# rpm0, rpm1, de, dr = control*(Fossen.control_ub - Fossen.control_lb) - Fossen.control_lb
# common subexpression elimination
x0 = np.cos(psi)
x1 = np.cos(theta)
x2 = u * x1
x3 = np.sin(phi)
x4 = np.sin(psi)
x5 = x3 * x4
x6 = np.sin(theta)
x7 = np.cos(phi)
x8 = x0 * x7
x9 = x4 * x7
x10 = x0 * x3
x11 = x1 * x3
x12 = x1 * x7
x13 = np.tan(theta)
x14 = q * x3
x15 = r * x7
x16 = 1 / x1
x17 = Yvv * abs(v)
x18 = m * xg
x19 = p * x18
x20 = m * zg
x21 = r * x20
x22 = m * u
x23 = m * yg
x24 = r * x23
x25 = np.sin(dr)
x26 = Kt0 * rpm0
x27 = Kt1 * rpm1
x28 = m * w
x29 = -x28
x30 = p * x23
x31 = -B + W
x32 = x3**2
x33 = x6**2
x34 = x1**2
x35 = x7**2
x36 = 1 / (x32 * x33 + x32 * x34 + x33 * x35 + x34 * x35)
x37 = x11 * x36
x38 = -p * (x29 - x30) - q * (x19 + x21) - r * (
x22 - x24) - v * x17 + x25 * (-x26 - x27) - x31 * x37
x39 = m**2
x40 = xg**4
x41 = x39 * x40
x42 = yg**2
x43 = xg**2
x44 = x39 * x43
x45 = x42 * x44
x46 = zg**2
x47 = x44 * x46
x48 = Iyy * Izz
x49 = Ixx * x48
x50 = Ixx * Izz
x51 = m * x50
x52 = m * x48
x53 = -x42 * x52 - x43 * x51 + x49
x54 = Ixx * Iyy
x55 = m * x54
x56 = -x43 * x55 - x46 * x52
x57 = Ixx * x41 + Iyy * x45 + Izz * x47 + x53 + x56
x58 = yg**4
x59 = x39 * x58
x60 = x42 * x46
x61 = x39 * x60
x62 = -x42 * x55 - x46 * x51
x63 = Ixx * x45 + Iyy * x59 + Izz * x61 + x62
x64 = zg**4
x65 = x39 * x64
x66 = Ixx * x47 + Iyy * x61 + Izz * x65
x67 = 1 / (x57 + x63 + x66)
x68 = x54 * yg
x69 = xg**3
x70 = Ixx * x69
x71 = yg**3
x72 = Iyy * m
x73 = x71 * x72
x74 = Izz * x46
x75 = x23 * x74
x76 = x67 * (x23 * x70 - x68 * xg + x73 * xg + x75 * xg)
x77 = x50 * zg
x78 = zg**3
x79 = Izz * x78
x80 = x20 * x42
x81 = Iyy * x80
x82 = x18 * x79 + x20 * x70 - x77 * xg + x81 * xg
x83 = Zww * abs(w)
x84 = m * v
x85 = p * x20
x86 = q * x23
x87 = np.sin(de)
x88 = np.cos(dr)
x89 = x88 * (x26 + x27)
x90 = -x22
x91 = q * x20
x92 = x12 * x36
x93 = -p * (x84 - x85) - q * (x90 - x91) - r * (
x19 + x86) - w * x83 - x31 * x92 + x87 * x89
x94 = x67 * x93
x95 = Xuu * abs(u)
x96 = q * x18
x97 = np.cos(de)
x98 = -x84
x99 = r * x18
x100 = x6 / (x33 + x34)
x101 = -p * (x21 + x86) - q * (x28 - x96) - r * (
x98 - x99) - u * x95 + x100 * x31 + x89 * x97
x102 = m**3
x103 = Ixx * x102
x104 = Iyy * x102
x105 = Izz * x102
x106 = x39 * x42
x107 = x39 * x46
x108 = x103 * x43
x109 = 1 / (m * x49 + x103 * x40 + x104 * x42 * x43 + x104 * x58 +
x104 * x60 + x105 * x43 * x46 + x105 * x60 + x105 * x64 -
x106 * x48 - x106 * x54 - x107 * x48 - x107 * x50 +
x108 * x42 + x108 * x46 - x44 * x50 - x44 * x54)
x110 = Ixx * x43
x111 = x110 * x23
x112 = Iyy * x23
x113 = -x111 - x112 * x46 + x68 - x73
x114 = p * q
x115 = Qt0 * rpm0
x116 = Qt1 * rpm1
x117 = x88 * (x115 + x116)
x118 = -x19
x119 = -x86
x120 = x67 * (B * (x100 * yb + x37 * xb) + Ixx * x114 - Iyy * x114 -
Nrr * r * abs(r) - W * (x100 * yg + x37 * xg) - u *
(x84 - x95 * ycp + x99) - v *
(x17 * xcp + x24 + x90) - w * (x118 + x119) + x117 * x87)
x121 = Izz * m
x122 = x121 * x78
x123 = Izz * x80 + x110 * x20 + x122 - x77
x124 = p * r
x125 = -x21
x126 = x67 * (B * (-x100 * zb - x92 * xb) - Ixx * x124 + Izz * x124 -
Mqq * q * abs(q) - W * (-x100 * zg - x92 * xg) - u *
(x29 + x95 * zcp + x96) - v * (x118 + x125) - w *
(x22 - x83 * xcp + x91) + x25 * (-x115 - x116))
x127 = x23 * xg
x128 = Izz * x127
x129 = x128 * zg
x130 = x112 * xg * zg
x131 = x129 - x130
x132 = q * r
x133 = x67 * (B * (-x37 * zb + x92 * yb) + Iyy * x132 - Izz * x132 -
Kpp * p * abs(p) - W * (-x37 * zg + x92 * yg) - u *
(x119 + x125) - v * (-x17 * zcp + x28 + x30) - w *
(x83 * ycp + x85 + x98) + x117 * x97)
x134 = x48 * zg
x135 = Iyy * x20
x136 = x111 * zg - x134 * yg + x135 * x71 + x23 * x79
x137 = Ixx * zg
x138 = x127 * x137
x139 = -x129 + x138
x140 = Ixx * m
x141 = x140 * x69
x142 = Ixx * x18
x143 = x42 * x72
x144 = x141 + x142 * x46 + x143 * xg - x54 * xg
x145 = -Izz * x20 * x43 - x122 + x134 - x81
x146 = x101 * x67
x147 = x38 * x67
x148 = x130 - x138
x149 = -x141 - x142 * x42 - x18 * x74 + x50 * xg
x150 = x112 * x43 - x48 * yg + x73 + x75
x151 = x39 * xg
x152 = x39 * x69
x153 = -x135 * xg + x151 * x42 * zg + x151 * x78 + x152 * zg
x154 = -x128 + x151 * x46 * yg + x151 * x71 + x152 * yg
x155 = -x121 * x46 + x45
x156 = -x143 + x47
x157 = -x137 * x23 + x39 * x71 * zg + x39 * x78 * yg + x44 * yg * zg
x158 = -x140 * x43 + x61
# dxdt
return np.array([
v * (x10 * x6 - x9) + w * (x5 + x6 * x8) + x0 * x2,
v * (x5 * x6 + x8) + w * (-x10 + x6 * x9) + x2 * x4,
-u * x6 + v * x11 + w * x12, p + x13 * x14 + x13 * x15,
q * x7 - r * x3, x14 * x16 + x15 * x16, x101 * x109 * x57 +
x113 * x120 + x123 * x126 + x131 * x133 + x38 * x76 + x82 * x94,
x101 * x76 + x109 * x38 * (x53 + x63) + x120 * x144 + x126 * x139 +
x133 * x145 + x136 * x94, x109 * x93 * (x49 + x56 + x62 + x66) +
x120 * x148 + x126 * x149 + x133 * x150 + x136 * x147 + x146 * x82,
x120 * x153 + x126 * x154 + x131 * x146 + x133 *
(-x121 * x43 + x155 + x156 + x41 - x43 * x72 + x48) + x145 * x147 +
x150 * x94, x120 * x157 + x123 * x146 + x126 *
(-x121 * x42 - x140 * x42 + x155 + x158 + x50 + x59) +
x133 * x154 + x139 * x147 + x149 * x94, x113 * x146 + x120 *
(-x140 * x46 + x156 + x158 - x46 * x72 + x54 + x65) + x126 * x157 +
x133 * x153 + x144 * x147 + x148 * x94
])
def _isf(self, q):
return jnp.tan(jnp.pi / 2.0 - jnp.pi * q)
def _isf(self, q):
return np.tan(np.pi / 2.0 - np.pi * q)
def t1_invcdf(u):
z = jnp.tan(jnp.pi * (u - 0.5))
return z
def f(c, a):
a1, a2 = a
c1, c2 = c
b = np.sum(np.cos(a1)) * np.sum(np.tan(c2 * a2))
c = c1 * np.sin(np.sum(a1 * a2)), c2 * np.cos(np.sum(a1))
return c, b
def bezierParsec(x):
xVals = np.array([0.0 for x in range(2 * len(uVals) + 2 * len(uValsBack))])
yVals = np.array([0.0 for x in range(2 * len(uVals) + 2 * len(uValsBack))])
uuVals = np.array(
[0.0 for x in range(2 * len(uVals) + 2 * len(uValsBack))])
#x = [rle, b8, xt, yt, b15, dzte, betate, b0, b2, xc, yc, gammale, b17, zte, alphate]
rle = x[0]
b8 = x[1]
xt = x[2]
yt = x[3]
b15 = x[4]
dzte = x[5]
betate = x[6]
b0 = x[7]
b2 = x[8]
xc = x[9]
yc = x[10]
gammale = x[11]
b17 = x[12]
zte = x[13]
alphate = x[14]
for i in range(2 * len(uVals) + 2 * len(uValsBack)):
if i < len(uValsBack) or i > (
(2 * len(uVals) + len(uValsBack)) - 1): #trailing edge
if i < len(uValsBack):
u = 1 - uValsBack[i]
else:
u = uValsBack[(i - (2 * len(uVals) + len(uValsBack)))] + (
1.0 - uValsBack[-1])
x0 = xt
x1 = (7.0 * xt - 9.0 * b8**2 / (2.0 * rle)) / 4.0
#x1 = (b15-xt)*.33 + xt
x2 = 3.0 * xt - 15.0 * b8**2.0 / (4 * rle)
#x2 = (b15-xt)*.66 + xt
x3 = b15
x4 = 1
y0 = yt
y1 = yt
y2 = (yt + b8) / 2
y3 = dzte + (1 - b15) * np.tan(betate)
y4 = dzte
xCoord = x0 * (1 - u)**4 + 4 * x1 * u * (1 - u)**3 + 6 * x2 * (
u**2) * (1 - u)**2 + 4 * x3 * (u**3) * (1 - u) + x4 * u**4
y_t = y0 * (1 - u)**4 + 4 * y1 * u * (1 - u)**3 + 6 * y2 * (
u**2) * (1 - u)**2 + 4 * y3 * (u**3) * (1 - u) + y4 * u**4
else: #leading edge
if i < (len(uVals) + len(uValsBack)):
u = uVals[-(i - len(uValsBack)) - 1]
else:
u = uVals[i -
(len(uVals) + len(uValsBack))] + (1.0 - uVals[-1])
if yt < (2.0 * rle * xt / 3.0)**0.5:
b8bound = yt
else:
b8bound = (2.0 * rle * xt / 3.0)**0.5
b8h = b8
if b8h > b8bound:
b8h = b8bound
if b8h < 0.0:
b8h = 0.0
x0 = 0
x1 = 0
x2 = 3.0 * (b8h**2.0) / (2.0 * rle)
x3 = xt
y0 = 0
y1 = b8h
y2 = yt
y3 = yt
xCoord = x0 * (1 - u)**3 + 3 * x1 * u * (1 - u)**2 + 3 * x2 * (
u**2) * (1 - u) + x3 * u**3
y_t = y0 * (1 - u)**3 + 3 * y1 * u * (1 - u)**2 + 3 * y2 * (
u**2) * (1 - u) + y3 * u**3
if xCoord < xc:
x0 = 0
x1 = b0
x2 = b2
x3 = xc
y0 = 0
y1 = b0 * np.tan(gammale)
y2 = yc
y3 = yc
#iterative search for u
u = (xCoord) / (xc)
for j in range(3):
f_x = x0 * (1 - u)**3 + 3 * x1 * u * (1 - u)**2 + 3 * x2 * (
u**2) * (1 - u) + x3 * u**3 - xCoord
fp_x = -3 * x0 * (1 -
u)**2 + 3 * x1 * (1 - u)**2 - 6 * x1 * u * (
1 - u) - 3 * x2 * u**2 + 6 * x2 * u * (
1 - u) + 3 * x3 * u**2
u = u - f_x / fp_x
y_c = y0 * (1 - u)**3 + 3 * y1 * u * (1 - u)**2 + 3 * y2 * (
u**2) * (1 - u) + y3 * u**3
else:
x0 = xc
x1 = (3.0 * xc - yc / np.tan(gammale)) / 2.0
x2 = (-8.0 * yc / np.tan(gammale) + 13.0 * xc) / 6.0
# x1 = .7 #used for debugging
# x2 = .9
x3 = b17
x4 = 1.0
y0 = yc
y1 = yc
y2 = 5.0 * yc / 6.0
y3 = zte - (1 - b17) * np.tan(alphate)
y4 = zte
#iterative search for u
u = (xCoord - xc) / (1.0 - xc)
for j in range(3):
f_x = x0 * (1 - u)**4 + 4 * x1 * u * (1 - u)**3 + 6 * x2 * (
u**2) * (1 - u)**2 + 4 * x3 * (u**3) * (
1 - u) + x4 * u**4 - xCoord
fp_x = -4 * x0 * (1 - u)**3 + 4 * x1 * (
1 - u)**3 - 12 * x1 * u * (1 - u)**2 + 12 * x2 * u * (
1 - u)**2 - 12 * x2 * (1 - u) * u**2 + 12 * x3 * (
1 - u) * u**2 - 4 * x3 * u**3 + 4 * x4 * u**3
if u < 0.25:
u = u - 1.0 * f_x / fp_x
else:
u = u - 1.0 * f_x / fp_x
y_c = y0 * (1 - u)**4 + 4 * y1 * u * (1 - u)**3 + 6 * y2 * (
u**2) * (1 - u)**2 + 4 * y3 * (u**3) * (1 - u) + y4 * u**4
if i < (len(uVals) + len(uValsBack)):
yVals = jo.index_update(yVals, i, y_c - y_t)
else:
yVals = jo.index_update(yVals, i, y_c + y_t)
xVals = jo.index_update(xVals, i, xCoord)
# uuVals = jo.index_update(uuVals, i, u) #used for debugging
return xVals, yVals
def loader(data_dir, filter_chain_options, device):
"""
Loads images from disk into a big numpy array, which will
later be pmapped onto all devices.
"""
splits = [entry for entry in data_dir.iterdir() if entry.is_dir()]
metadata = {
split: json.load(
StringIO((data_dir / f"transforms_{split.name}.json").read_text())
)
for split in splits
}
vmap_filter_chain = vmap(
lambda imgs: jit(filter_chain, static_argnums=(1,), device=device)(
imgs, filter_chain_options
),
)
frame_iterator = lambda f, mdata: jnp.stack(
[
f(frame)
for idx, frame in enumerate(mdata["frames"])
if idx % filter_chain_options.skiptest == 0
],
axis=0,
)
images = {
split.name: vmap_filter_chain(
jnp.array(
frame_iterator(
lambda frame: imageio.imread(
data_dir / f"{frame['file_path']}.png"
),
mdata,
)
)
)
for split, mdata in metadata.items()
}
poses = {
split.name: frame_iterator(
lambda frame: jnp.array(frame["transform_matrix"]), mdata
)
for split, mdata in metadata.items()
}
intrinsics = {
split.name: Intrinsics(
focal_length=0.5
* images[split.name].shape[2]
/ jnp.tan(0.5 * float(mdata["camera_angle_x"])),
width=images[split.name].shape[2],
height=images[split.name].shape[1],
)
for split, mdata in metadata.items()
}
return images, poses, intrinsics
def _ppf(self, q):
return np.tan(np.pi * q - np.pi / 2.0)
softsign = utils.copy_docstring(
tf.math.softsign,
lambda features, name=None: np.divide(features, (np.abs(features) + 1)))
sqrt = utils.copy_docstring(tf.math.sqrt, lambda x, name=None: np.sqrt(x))
square = utils.copy_docstring(tf.math.square,
lambda x, name=None: np.square(x))
squared_difference = utils.copy_docstring(
tf.math.squared_difference, lambda x, y, name=None: np.square(x - y))
subtract = utils.copy_docstring(tf.math.subtract,
lambda x, y, name=None: np.subtract(x, y))
tan = utils.copy_docstring(tf.math.tan, lambda x, name=None: np.tan(x))
tanh = utils.copy_docstring(tf.math.tanh, lambda x, name=None: np.tanh(x))
top_k = utils.copy_docstring(tf.math.top_k, _top_k)
truediv = utils.copy_docstring(tf.math.truediv,
lambda x, y, name=None: np.true_divide(x, y))
# unsorted_segment_max = utils.copy_docstring(
# tf.math.unsorted_segment_max,
# lambda data, segment_ids, num_segments, name=None: (
# np.unsorted_segment_max))
# unsorted_segment_mean = utils.copy_docstring(
# tf.math.unsorted_segment_mean,
def g(z): return 3. * np.exp(np.sin(x).sum() * np.cos(y) * np.tan(z))
def __lsp(A, phi, theta):
return (A * jnp.exp(theta * jnp.tan(jnp.deg2rad(phi))) * jnp.stack(
(jnp.cos(theta), jnp.sin(theta)))).T
def correct_logsp_params(A, phi, q, psi, dpsi, theta):
Ap = jnp.exp(-dpsi * jnp.tan(jnp.deg2rad(phi)))
return A * Ap, phi, q, psi, theta + dpsi
def standard_cauchy_icdf(u):
"""Inverse cumulative distribution function of standard Cauchy."""
return np.tan(np.pi * (u - 0.5))