from functools import partial

import jax.numpy as jnp

import numpy as np

import tensorflow as tf

import torch

from jax import lax

def tridiagonal_solve(a: np.ndarray, b: np.ndarray, c: np.ndarray, d: np.ndarray) -> np.ndarray:

n = len(b)

# Create working copies to avoid modifying input

c_prime = np.zeros(n - 1, dtype=np.float64)

d_prime = np.zeros(n, dtype=np.float64)

x = np.zeros(n, dtype=np.float64)

# Forward sweep - sequential dependency: c_prime[i] depends on c_prime[i-1]

c_prime[0] = c[0] / b[0]

d_prime[0] = d[0] / b[0]

for i in range(1, n - 1):

denom = b[i] - a[i - 1] * c_prime[i - 1]

c_prime[i] = c[i] / denom

d_prime[i] = (d[i] - a[i - 1] * d_prime[i - 1]) / denom

# Last row of forward sweep

denom = b[n - 1] - a[n - 2] * c_prime[n - 2]

d_prime[n - 1] = (d[n - 1] - a[n - 2] * d_prime[n - 2]) / denom

# Back substitution - sequential dependency: x[i] depends on x[i+1]

x[n - 1] = d_prime[n - 1]

for i in range(n - 2, -1, -1):

x[i] = d_prime[i] - c_prime[i] * x[i + 1]

return x

def leapfrog_integration(

positions: np.ndarray,

velocities: np.ndarray,

masses: np.ndarray,

dt: float,

n_steps: int,

softening: float = 0.01

) -> tuple[np.ndarray, np.ndarray]:

n_particles = len(masses)

pos = positions.copy()

vel = velocities.copy()

acc = np.zeros_like(pos)

G = 1.0

for step in range(n_steps):

acc.fill(0.0)

for i in range(n_particles):

for j in range(i + 1, n_particles):

dx = pos[j, 0] - pos[i, 0]

dy = pos[j, 1] - pos[i, 1]

dz = pos[j, 2] - pos[i, 2]

dist_sq = dx * dx + dy * dy + dz * dz + softening * softening

dist = np.sqrt(dist_sq)

dist_cubed = dist_sq * dist

force_over_dist = G / dist_cubed

acc[i, 0] += masses[j] * force_over_dist * dx

acc[i, 1] += masses[j] * force_over_dist * dy

acc[i, 2] += masses[j] * force_over_dist * dz

acc[j, 0] -= masses[i] * force_over_dist * dx

acc[j, 1] -= masses[i] * force_over_dist * dy

acc[j, 2] -= masses[i] * force_over_dist * dz

for i in range(n_particles):

vel[i, 0] += 0.5 * dt * acc[i, 0]

vel[i, 1] += 0.5 * dt * acc[i, 1]

vel[i, 2] += 0.5 * dt * acc[i, 2]

for i in range(n_particles):

pos[i, 0] += dt * vel[i, 0]

pos[i, 1] += dt * vel[i, 1]

pos[i, 2] += dt * vel[i, 2]

for i in range(n_particles):

vel[i, 0] += 0.5 * dt * acc[i, 0]

vel[i, 1] += 0.5 * dt * acc[i, 1]

vel[i, 2] += 0.5 * dt * acc[i, 2]

return pos, vel

def longest_increasing_subsequence_length(arr: np.ndarray) -> int:

n = len(arr)

if n == 0:

return 0

dp = np.ones(n, dtype=np.int64)

for i in range(1, n):

for j in range(i):

if arr[j] < arr[i]:

if dp[j] + 1 > dp[i]:

dp[i] = dp[j] + 1

max_length = dp[0]

for i in range(1, n):

if dp[i] > max_length:

max_length = dp[i]

return max_length

def _tridiagonal_forward_step_jax(carry, inputs):

c_prev, d_prev = carry

a_i, b_i, c_i, d_i = inputs

denom = b_i - a_i * c_prev

c_new = c_i / denom

d_new = (d_i - a_i * d_prev) / denom

return (c_new, d_new), (c_new, d_new)

def _tridiagonal_back_step_jax(x_next, inputs):

d_prime_i, c_prime_i = inputs

x_i = d_prime_i - c_prime_i * x_next

return x_i, x_i

def tridiagonal_solve_jax(a, b, c, d):

n = b.shape[0]

c_prime_0 = c[0] / b[0]

d_prime_0 = d[0] / b[0]

scan_inputs = (a[:-1], b[1:-1], c[1:], d[1:-1])

_, (c_prime_rest, d_prime_mid) = lax.scan(

_tridiagonal_forward_step_jax,

(c_prime_0, d_prime_0),

scan_inputs

)

c_prime = jnp.concatenate([jnp.array([c_prime_0]), c_prime_rest])

denom_last = b[n - 1] - a[n - 2] * c_prime[n - 2]

d_prime_last = (d[n - 1] - a[n - 2] * d_prime_mid[-1]) / denom_last

d_prime = jnp.concatenate([jnp.array([d_prime_0]), d_prime_mid, jnp.array([d_prime_last])])

x_last = d_prime[n - 1]

_, x_rest = lax.scan(

_tridiagonal_back_step_jax,

x_last,

(d_prime[:-1], c_prime),

reverse=True

)

x = jnp.concatenate([x_rest, jnp.array([x_last])])

return x

def _leapfrog_compute_accelerations_jax(pos, masses, softening):

G = 1.0

diff = pos[jnp.newaxis, :, :] - pos[:, jnp.newaxis, :]

dist_sq = jnp.sum(diff ** 2, axis=-1) + softening ** 2

dist = jnp.sqrt(dist_sq)

dist_cubed = dist_sq * dist

dist_cubed = jnp.where(dist_cubed == 0, 1.0, dist_cubed)

force_factor = G * masses[jnp.newaxis, :] / dist_cubed

acc = jnp.sum(force_factor[:, :, jnp.newaxis] * diff, axis=1)

return acc

def _leapfrog_step_jax(carry, _, masses, softening, dt):

pos, vel = carry

acc = _leapfrog_compute_accelerations_jax(pos, masses, softening)

vel = vel + 0.5 * dt * acc

pos = pos + dt * vel

vel = vel + 0.5 * dt * acc

return (pos, vel), None

def leapfrog_integration_jax(

positions,

velocities,

masses,

dt: float,

n_steps: int,

softening: float = 0.01

):

step_fn = partial(_leapfrog_step_jax, masses=masses, softening=softening, dt=dt)

(final_pos, final_vel), _ = lax.scan(step_fn, (positions, velocities), None, length=n_steps)

return final_pos, final_vel

def _lis_inner_body_jax(j, dp_inner, arr, i):

condition = (arr[j] < arr[i]) & (dp_inner[j] + 1 > dp_inner[i])

new_val = jnp.where(condition, dp_inner[j] + 1, dp_inner[i])

return dp_inner.at[i].set(new_val)

def _lis_outer_body_jax(i, dp, arr):

inner_fn = partial(_lis_inner_body_jax, arr=arr, i=i)

dp = lax.fori_loop(0, i, inner_fn, dp)

return dp

def longest_increasing_subsequence_length_jax(arr):

n = arr.shape[0]

if n == 0:

return 0

outer_fn = partial(_lis_outer_body_jax, arr=arr)

dp = jnp.ones(n, dtype=jnp.int32)

dp = lax.fori_loop(1, n, outer_fn, dp)

return int(jnp.max(dp))

def tridiagonal_solve_torch(a, b, c, d):

device = b.device

dtype = b.dtype

n = b.shape[0]

c_prime = torch.zeros(n - 1, device=device, dtype=dtype)

d_prime = torch.zeros(n, device=device, dtype=dtype)

x = torch.zeros(n, device=device, dtype=dtype)

c_prime[0] = c[0] / b[0]

d_prime[0] = d[0] / b[0]

for i in range(1, n - 1):

denom = b[i] - a[i - 1] * c_prime[i - 1]

c_prime[i] = c[i] / denom

d_prime[i] = (d[i] - a[i - 1] * d_prime[i - 1]) / denom

denom = b[n - 1] - a[n - 2] * c_prime[n - 2]

d_prime[n - 1] = (d[n - 1] - a[n - 2] * d_prime[n - 2]) / denom

x[n - 1] = d_prime[n - 1]

for i in range(n - 2, -1, -1):

x[i] = d_prime[i] - c_prime[i] * x[i + 1]

return x

def leapfrog_integration_torch(

positions,

velocities,

masses,

dt: float,

n_steps: int,

softening: float = 0.01

):

G = 1.0

pos = positions.clone()

vel = velocities.clone()

for _ in range(n_steps):

diff = pos.unsqueeze(0) - pos.unsqueeze(1)

dist_sq = torch.sum(diff ** 2, dim=-1) + softening ** 2

dist = torch.sqrt(dist_sq)

dist_cubed = dist_sq * dist

dist_cubed = torch.where(dist_cubed == 0, torch.ones_like(dist_cubed), dist_cubed)

force_factor = G * masses.unsqueeze(0) / dist_cubed

acc = torch.sum(force_factor.unsqueeze(-1) * diff, dim=1)

vel = vel + 0.5 * dt * acc

pos = pos + dt * vel

vel = vel + 0.5 * dt * acc

return pos, vel

def longest_increasing_subsequence_length_torch(arr):

n = arr.shape[0]

if n == 0:

return 0

device = arr.device

dp = torch.ones(n, device=device, dtype=torch.int64)

for i in range(1, n):

for j in range(i):

if arr[j] < arr[i]:

if dp[j] + 1 > dp[i]:

dp[i] = dp[j] + 1

return int(torch.max(dp).item())

def _tridiagonal_forward_cond_tf(i, _c_prime, _d_prime, n, _a, _b, _c, _d):

return i < n - 1

def _tridiagonal_forward_body_tf(i, c_prime, d_prime, n, a, b, c, d):

c_prev = c_prime[i - 1]

d_prev = d_prime[i - 1]

denom = b[i] - a[i - 1] * c_prev

c_val = c[i] / denom

d_val = (d[i] - a[i - 1] * d_prev) / denom

c_prime = tf.tensor_scatter_nd_update(c_prime, tf.reshape(i, [1, 1]), tf.reshape(c_val, [1]))

d_prime = tf.tensor_scatter_nd_update(d_prime, tf.reshape(i, [1, 1]), tf.reshape(d_val, [1]))

return i + 1, c_prime, d_prime, n, a, b, c, d

def _tridiagonal_back_cond_tf(i, _x, _c_prime, _d_prime):

return i >= 0

def _tridiagonal_back_body_tf(i, x, c_prime, d_prime):

x_next = x[i + 1]

x_val = d_prime[i] - c_prime[i] * x_next

x = tf.tensor_scatter_nd_update(x, tf.reshape(i, [1, 1]), tf.reshape(x_val, [1]))

return i - 1, x, c_prime, d_prime

def tridiagonal_solve_tf(a, b, c, d):

n = tf.shape(b)[0]

dtype = b.dtype

c_prime = tf.zeros([n - 1], dtype=dtype)

d_prime = tf.zeros([n], dtype=dtype)

c_prime = tf.tensor_scatter_nd_update(c_prime, [[0]], tf.reshape(c[0] / b[0], [1]))

d_prime = tf.tensor_scatter_nd_update(d_prime, [[0]], tf.reshape(d[0] / b[0], [1]))

_, c_prime, d_prime, _, _, _, _, _ = tf.while_loop(

_tridiagonal_forward_cond_tf,

_tridiagonal_forward_body_tf,

[1, c_prime, d_prime, n, a, b, c, d]

)

c_last = c_prime[n - 2]

d_prev = d_prime[n - 2]

denom = b[n - 1] - a[n - 2] * c_last

d_last = (d[n - 1] - a[n - 2] * d_prev) / denom

d_prime = tf.tensor_scatter_nd_update(d_prime, tf.reshape(n - 1, [1, 1]), tf.reshape(d_last, [1]))

x = tf.zeros([n], dtype=dtype)

x = tf.tensor_scatter_nd_update(x, tf.reshape(n - 1, [1, 1]), tf.reshape(d_prime[n - 1], [1]))

_, x, _, _ = tf.while_loop(

_tridiagonal_back_cond_tf,

_tridiagonal_back_body_tf,

[n - 2, x, c_prime, d_prime]

)

return x

def _leapfrog_compute_accelerations_tf(pos, masses, softening, G):

diff = tf.expand_dims(pos, 0) - tf.expand_dims(pos, 1)

dist_sq = tf.reduce_sum(diff ** 2, axis=-1) + softening ** 2

dist = tf.sqrt(dist_sq)

dist_cubed = dist_sq * dist

dist_cubed = tf.where(dist_cubed == 0, tf.ones_like(dist_cubed), dist_cubed)

force_factor = G * tf.expand_dims(masses, 0) / dist_cubed

acc = tf.reduce_sum(tf.expand_dims(force_factor, -1) * diff, axis=1)

return acc

def _leapfrog_step_body_tf(i, pos, vel, masses, softening, dt, n_steps):

G = 1.0

acc = _leapfrog_compute_accelerations_tf(pos, masses, softening, G)

vel = vel + 0.5 * dt * acc

pos = pos + dt * vel

vel = vel + 0.5 * dt * acc

return i + 1, pos, vel, masses, softening, dt, n_steps

def _leapfrog_step_cond_tf(i, _pos, _vel, _masses, _softening, _dt, n_steps):

return i < n_steps

def leapfrog_integration_tf(

positions,

velocities,

masses,

dt: float,

n_steps: int,

softening: float = 0.01

):

dt = tf.constant(dt, dtype=positions.dtype)

softening = tf.constant(softening, dtype=positions.dtype)

_, final_pos, final_vel, _, _, _, _ = tf.while_loop(

_leapfrog_step_cond_tf,

_leapfrog_step_body_tf,

[0, positions, velocities, masses, softening, dt, n_steps]

)

return final_pos, final_vel

def _lis_inner_body_tf(j, dp_inner, arr, i):

condition = tf.logical_and(arr[j] < arr[i], dp_inner[j] + 1 > dp_inner[i])

new_val = tf.where(condition, dp_inner[j] + 1, dp_inner[i])

indices = tf.reshape(i, [1, 1])

updates = tf.reshape(new_val, [1])

dp_updated = tf.tensor_scatter_nd_update(dp_inner, indices, updates)

return j + 1, dp_updated, arr, i

def _lis_inner_cond_tf(j, _dp_inner, _arr, i):

return j < i

def _lis_outer_body_tf(i, dp, arr, n):

_, dp, _, _ = tf.while_loop(

_lis_inner_cond_tf,

_lis_inner_body_tf,

[0, dp, arr, i]

)

return i + 1, dp, arr, n

def _lis_outer_cond_tf(i, _dp, _arr, n):

return i < n

def longest_increasing_subsequence_length_tf(arr):

n = tf.shape(arr)[0]

if n == 0:

return 0

dp = tf.ones(n, dtype=tf.int32)

_, dp, _, _ = tf.while_loop(

_lis_outer_cond_tf,

_lis_outer_body_tf,

[1, dp, arr, n]

)

return int(tf.reduce_max(dp))