Skip to content

OPT: precompute radar geometry and point weights #170

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rcjackson merged 1 commit into from
Oct 2, 2026
Merged

OPT: precompute radar geometry and point weights #170

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 40 additions & 26 deletions pydda/cost_functions/_cost_functions_jax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ def _apply_upper_bc(grad_w, upper_bc, upper_bc_mask=None):


def calculate_radial_vel_cost_function(
vrs, azs, els, u, v, w, wts, rmsVr, weights, coeff=1.0
vrs, azs, els, u, v, w, wts, rmsVr, weights, coeff=1.0, radar_geometry=None
):
"""
Calculates the cost function due to difference of the wind field from
Expand Down Expand Up @@ -80,11 +80,17 @@ def calculate_radial_vel_cost_function(
J_o = 0
lambda_o = coeff / (rmsVr * rmsVr)
for i in range(len(vrs)):
v_ar = (
jnp.cos(els[i]) * jnp.sin(azs[i]) * u
+ jnp.cos(els[i]) * jnp.cos(azs[i]) * v
+ jnp.sin(els[i]) * (w - jnp.abs(wts[i]))
)
if radar_geometry is None:
x_projection = jnp.cos(els[i]) * jnp.sin(azs[i])
y_projection = jnp.cos(els[i]) * jnp.cos(azs[i])
sin_el = jnp.sin(els[i])
wts_abs = jnp.abs(wts[i])
else:
x_projection = radar_geometry["x_projection"][i]
y_projection = radar_geometry["y_projection"][i]
sin_el = radar_geometry["sin_el"][i]
wts_abs = radar_geometry["wts_abs"][i]
v_ar = x_projection * u + y_projection * v + sin_el * (w - wts_abs)
the_weight = jnp.asarray(weights[i])
J_o += lambda_o * jnp.sum(jnp.square(vrs[i] - v_ar) * the_weight)
return J_o
Expand Down Expand Up @@ -307,7 +313,7 @@ def calculate_smoothness_gradient(
return y.flatten()


def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3):
def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, point_weights=None):
"""
Calculates the cost function related to point observations. A mean square error cost
function term is applied to points that are within the sphere of influence
Expand Down Expand Up @@ -344,21 +350,26 @@ def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3):
wind field and points.
"""
J = 0.0
for the_point in point_list:
dist = jnp.sqrt(
(x - the_point["x"]) ** 2
+ (y - the_point["y"]) ** 2
+ (z - the_point["z"]) ** 2
)
dist = jnp.maximum(dist, 1.0)
weight = 1 / dist**2
weight = weight / jnp.sum(weight)
for point_index, the_point in enumerate(point_list):
if point_weights is None:
dist = jnp.sqrt(
(x - the_point["x"]) ** 2
+ (y - the_point["y"]) ** 2
+ (z - the_point["z"]) ** 2
)
dist = jnp.maximum(dist, 1.0)
weight = 1 / dist**2
weight = weight / jnp.sum(weight)
else:
weight = point_weights[point_index]
J += jnp.sum(weight * ((u - the_point["u"]) ** 2 + (v - the_point["v"]) ** 2))

return J * Cp


def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0):
def calculate_point_gradient(
u, v, x, y, z, point_list, Cp=1e-3, roi=500.0, point_weights=None
):
"""
Calculates the gradient of the cost function related to point observations.
A mean square error cost function term is applied to points that are within the sphere of influence
Expand Down Expand Up @@ -398,15 +409,18 @@ def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0):
gradJ_v = jnp.zeros_like(v)
gradJ_w = jnp.zeros_like(u)

for the_point in point_list:
dist = jnp.sqrt(
(x - the_point["x"]) ** 2
+ (y - the_point["y"]) ** 2
+ (z - the_point["z"]) ** 2
)
dist = jnp.maximum(dist, 1.0)
weight = 1 / dist**2
weight = weight / jnp.sum(weight)
for point_index, the_point in enumerate(point_list):
if point_weights is None:
dist = jnp.sqrt(
(x - the_point["x"]) ** 2
+ (y - the_point["y"]) ** 2
+ (z - the_point["z"]) ** 2
)
dist = jnp.maximum(dist, 1.0)
weight = 1 / dist**2
weight = weight / jnp.sum(weight)
else:
weight = point_weights[point_index]
gradJ_u += 2 * (u - the_point["u"]) * weight
gradJ_v += 2 * (v - the_point["v"]) * weight

Expand Down
69 changes: 49 additions & 20 deletions pydda/cost_functions/_cost_functions_numpy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -249,6 +249,7 @@ def calculate_radial_vel_cost_and_gradient(
upper_bc=1,
upper_bc_mask=None,
parallel=False,
radar_geometry=None,
):
"""Calculate radar cost and gradient while sharing the radar projection.

Expand All @@ -260,17 +261,24 @@ def calculate_radial_vel_cost_and_gradient(

if parallel:
vrs_arr = np.stack(vrs)
els_arr = np.stack(els)
azs_arr = np.stack(azs)
wts_arr = np.stack(wts)
cos_els = np.cos(els_arr)
sin_els = np.sin(els_arr)
x_projection = cos_els * np.sin(azs_arr)
y_projection = cos_els * np.cos(azs_arr)
if radar_geometry is None:
els_arr = np.stack(els)
azs_arr = np.stack(azs)
wts_arr = np.stack(wts)
cos_els = np.cos(els_arr)
sin_els = np.sin(els_arr)
x_projection = cos_els * np.sin(azs_arr)
y_projection = cos_els * np.cos(azs_arr)
wts_abs = np.abs(wts_arr)
else:
x_projection = radar_geometry["x_projection"]
y_projection = radar_geometry["y_projection"]
sin_els = radar_geometry["sin_el"]
wts_abs = radar_geometry["wts_abs"]
v_ar = (
x_projection * u[np.newaxis]
+ y_projection * v[np.newaxis]
+ sin_els * (w[np.newaxis] - np.abs(wts_arr))
+ sin_els * (w[np.newaxis] - wts_abs)
)
residual = v_ar - vrs_arr
cost = lambda_o * np.sum(np.square(residual) * weights)
Expand All @@ -282,12 +290,21 @@ def calculate_radial_vel_cost_and_gradient(
p_x1 = np.zeros(vrs[0].shape)
p_y1 = np.zeros(vrs[0].shape)
p_z1 = np.zeros(vrs[0].shape)
for vr, el, az, wt, weight in zip(vrs, els, azs, wts, weights):
cos_el = np.cos(el)
sin_el = np.sin(el)
x_projection = cos_el * np.sin(az)
y_projection = cos_el * np.cos(az)
v_ar = x_projection * u + y_projection * v + sin_el * (w - np.abs(wt))
for index, (vr, wt, weight) in enumerate(zip(vrs, wts, weights)):
if radar_geometry is None:
el = els[index]
az = azs[index]
cos_el = np.cos(el)
sin_el = np.sin(el)
x_projection = cos_el * np.sin(az)
y_projection = cos_el * np.cos(az)
wts_abs = np.abs(wt)
else:
x_projection = radar_geometry["x_projection"][index]
y_projection = radar_geometry["y_projection"][index]
sin_el = radar_geometry["sin_el"][index]
wts_abs = radar_geometry["wts_abs"][index]
v_ar = x_projection * u + y_projection * v + sin_el * (w - wts_abs)
residual = v_ar - vr
cost += lambda_o * np.sum(np.square(residual) * weight)
p_x1 += 2 * residual * x_projection * weight * lambda_o
Expand Down Expand Up @@ -457,7 +474,9 @@ def _point_weights(x, y, z, the_point, roi):
return weight / maximum if maximum > 0 else weight


def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, power=2, roi=500.0):
def calculate_point_cost(
u, v, x, y, z, point_list, Cp=1e-3, power=2, roi=500.0, point_weights=None
):
"""
Calculates the cost function related to point observations. A mean square error cost
function term is applied to points that are within the sphere of influence
Expand Down Expand Up @@ -490,18 +509,24 @@ def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, power=2, roi=500.0)
The cost function related to the difference between wind field and points.
"""
J = 0.0
for the_point in point_list:
for point_index, the_point in enumerate(point_list):
# Instead of worrying about whole domain, just find points in radius of influence
# Since we know that the weight will be zero outside the sphere of influence anyways

weight = _point_weights(x, y, z, the_point, roi)
weight = (
point_weights[point_index]
if point_weights is not None
else _point_weights(x, y, z, the_point, roi)
)

J += np.sum(weight * ((u - the_point["u"]) ** 2 + (v - the_point["v"]) ** 2))

return J * Cp


def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0):
def calculate_point_gradient(
u, v, x, y, z, point_list, Cp=1e-3, roi=500.0, point_weights=None
):
"""
Calculates the gradient of the cost function related to point observations.
A mean square error cost function term is applied to points that are within the sphere of influence
Expand Down Expand Up @@ -537,8 +562,12 @@ def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0):
gradJ_v = np.zeros_like(v)
gradJ_w = np.zeros_like(u)

for the_point in point_list:
weight = _point_weights(x, y, z, the_point, roi)
for point_index, the_point in enumerate(point_list):
weight = (
point_weights[point_index]
if point_weights is not None
else _point_weights(x, y, z, the_point, roi)
)
gradJ_u += 2 * weight * (u - the_point["u"])
gradJ_v += 2 * weight * (v - the_point["v"])

Expand Down
74 changes: 45 additions & 29 deletions pydda/cost_functions/_cost_functions_tensorflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,7 @@ def calculate_radial_vel_cost_and_gradient(
upper_bc=1,
upper_bc_mask=None,
lower_bc=True,
radar_geometry=None,
):
"""Calculate radar cost and gradient while sharing each radar residual."""
lambda_o = coeff / (rmsVr * rmsVr)
Expand All @@ -118,17 +119,22 @@ def calculate_radial_vel_cost_and_gradient(
p_z1 = tf.zeros_like(w)

for i in range(len(vrs)):
cos_el = tf.math.cos(els[i])
sin_el = tf.math.sin(els[i])
v_ar = (
cos_el * tf.math.sin(azs[i]) * u
+ cos_el * tf.math.cos(azs[i]) * v
+ sin_el * (w - tf.math.abs(wts[i]))
)
if radar_geometry is None:
cos_el = tf.math.cos(els[i])
sin_el = tf.math.sin(els[i])
x_projection = cos_el * tf.math.sin(azs[i])
y_projection = cos_el * tf.math.cos(azs[i])
wts_abs = tf.math.abs(wts[i])
else:
x_projection = radar_geometry["x_projection"][i]
y_projection = radar_geometry["y_projection"][i]
sin_el = radar_geometry["sin_el"][i]
wts_abs = radar_geometry["wts_abs"][i]
v_ar = x_projection * u + y_projection * v + sin_el * (w - wts_abs)
residual = v_ar - vrs[i]
costs.append(lambda_o * tf.reduce_sum(tf.math.square(residual) * weights[i]))
p_x1 += 2 * residual * cos_el * tf.math.sin(azs[i]) * weights[i] * lambda_o
p_y1 += 2 * residual * cos_el * tf.math.cos(azs[i]) * weights[i] * lambda_o
p_x1 += 2 * residual * x_projection * weights[i] * lambda_o
p_y1 += 2 * residual * y_projection * weights[i] * lambda_o
p_z1 += 2 * residual * sin_el * weights[i] * lambda_o

if lower_bc:
Expand Down Expand Up @@ -383,7 +389,9 @@ def calculate_smoothness_gradient(
return tf.reshape(y, (3 * np.prod(u.shape),))


def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0):
def calculate_point_cost(
u, v, x, y, z, point_list, Cp=1e-3, roi=500.0, point_weights=None
):
"""
Calculates the cost function related to point observations. A mean square error cost
function term is applied to points that are within the sphere of influence
Expand Down Expand Up @@ -421,26 +429,31 @@ def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0):
"""
J = tf.Variable(0.0)

for the_point in point_list:
for point_index, the_point in enumerate(point_list):
# Instead of worrying about whole domain, just find points in radius of influence
# Since we know that the weight will be zero outside the sphere of influence anyways
up = tf.ones_like(u) * the_point["u"]
vp = tf.ones_like(v) * the_point["v"]
dist = tf.math.sqrt(
(x - the_point["x"]) ** 2
+ (y - the_point["y"]) ** 2
+ (z - the_point["z"]) ** 2
)
dist = tf.math.maximum(dist, 1.0)
weight = 1 / dist**2
weight = weight / tf.reduce_max(weight)
if point_weights is None:
dist = tf.math.sqrt(
(x - the_point["x"]) ** 2
+ (y - the_point["y"]) ** 2
+ (z - the_point["z"]) ** 2
)
dist = tf.math.maximum(dist, 1.0)
weight = 1 / dist**2
weight = weight / tf.reduce_max(weight)
else:
weight = point_weights[point_index]

J.assign_add(tf.math.reduce_sum(((u - up) ** 2 + (v - vp) ** 2) * weight))

return J * Cp


def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0):
def calculate_point_gradient(
u, v, x, y, z, point_list, Cp=1e-3, roi=500.0, point_weights=None
):
"""
Calculates the gradient of the cost function related to point observations.
A mean square error cost function term is applied to points that are within the sphere of influence
Expand Down Expand Up @@ -477,20 +490,23 @@ def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0):
gradJ_u = tf.Variable(tf.zeros(u.shape, dtype=tf.float32))
gradJ_v = tf.Variable(tf.zeros(v.shape, dtype=tf.float32))
gradJ_w = tf.Variable(tf.zeros(u.shape, dtype=tf.float32))
for the_point in point_list:
for point_index, the_point in enumerate(point_list):
# Instead of worrying about whole domain, just find points in radius of influence
# Since we know that the weight will be zero outside the sphere of influence anyways
up = tf.ones_like(u, dtype=tf.float32) * the_point["u"]
vp = tf.ones_like(v, dtype=tf.float32) * the_point["v"]

dist = tf.math.sqrt(
(x - the_point["x"]) ** 2
+ (y - the_point["y"]) ** 2
+ (z - the_point["z"]) ** 2
)
dist = tf.math.maximum(dist, 1.0)
weight = 1 / dist**2
weight = weight / tf.reduce_max(weight)
if point_weights is None:
dist = tf.math.sqrt(
(x - the_point["x"]) ** 2
+ (y - the_point["y"]) ** 2
+ (z - the_point["z"]) ** 2
)
dist = tf.math.maximum(dist, 1.0)
weight = 1 / dist**2
weight = weight / tf.reduce_max(weight)
else:
weight = point_weights[point_index]
gradJ_u.assign_add((2 * (u - up) * weight))
gradJ_v.assign_add((2 * (v - vp) * weight))
gradJ = tf.stack([gradJ_u, gradJ_v, gradJ_w], axis=0)
Expand Down
Loading
Loading