for loop to built-in array-oriented function to avoid accumulation of floating-point roundoff error.

[mjhoover1]

Edited:
Hi, I was wondering if it would be possible to transition such a for loop in this function to a built-in array-oriented function?
Complete Minimal Example:

import jax.numpy as jnp
from jax._src.api import value_and_grad
from jax import grad, jit, lax
# again, this only works on startup!
from jax.config import config
from jax.scipy.special import logsumexp
config.update('jax_platform_name', 'cpu')
config.update("jax_debug_nans", True)
config.update("jax_enable_x64", True)

@jit
def sigmoid(x):
    sigmoid_slope = 32.
    return 1.0 / (1.0 + jnp.exp(-sigmoid_slope * x))

@jit
def sigmoid2(x, x_offset):
    sigmoid_slope = 32.
    return 1. / (1. + jnp.exp(-sigmoid_slope * (x - x_offset)))

@jit
def in_range_unscaled(x, from_, to):
    return sigmoid2(x, from_)*(1. - sigmoid2(x, to))

@jit
def in_range(x, from_, to):
    return in_range_unscaled(x, from_, to) / in_range_unscaled((from_ + to) / 2.0, from_, to) # second term just for scaling

def min_pos_delta_if_equal(x, num_elems, ref_idx, y, z):
    sum = 0.0
    M = 1000.
    for i in range(num_elems):
        if i == ref_idx: continue
        diff = x[i] - x[ref_idx]
        neg_offset = 2 * M * (1.0 - sigmoid(diff))
        diff += neg_offset
        lane_shift = 1000. * (1.0 - in_range(y[i] - z, -0.5, 0.5))
        diff += lane_shift
        exp_arg = -diff
        sum += jnp.exp(exp_arg)

    return -jnp.log(sum)

if __name__ == "__main__":
    print(min_pos_delta_if_equal([0. 0. 0. 0.], 4, 0, [0. 0. 0. 0.], 0.0))

I can't seem to see how to incorporate the if condition, here is my attempt:

return -jnp.logsumexp( (x[:num_elems]-x[ref_idx]) + 
                                      (2*M*(1.0-sigmoid(x[:num_elems]-x[ref_idx]) + 
                                      (1000. * (1. - in_range(y[:num_elems] - z, -.5, .5)))
                                     )

If not would an alternative solution to avoid roundoff error be to use lax.foriloop. I attempted to transition the function using this, but it doesn't seem to work since i depends on a condition, thanks!

def min_pos_delta_if_equal(x, num_elems, ref_idx, y, z):
    sum = 0.0
    M = 1000.
    def body_fun(i, sum):
        if i != ref_idx:
            diff = x[i] - x[ref_idx]
            neg_offset = 2 * M * (1. - sigmoid(diff))
            diff += neg_offset
            lane_shift = 1000. * (1.0 - in_range(y[i] - z, -0.5, 0.5))
            diff += lane_shift
            exp_arg = -diff
            sum += jnp.exp(exp_arg)
            return sum
    sum = lax.fori_loop(0, num_elems, body_fun, 0)
    return -jnp.log(sum)

[jakevdp]

It looks like this could be done in a vectorized fashion, but I don't know what in_range is, and I'd have to guess at the shape of the inputs. Could you provide a complete example of what you want to compute (including imports and example inputs?)

[mjhoover1]

This was actually an attempt to distill the code to a few lines and is actually derived from the previous question #8734 to clarify here is a complete example:

from operator import ne, neg
import jax.numpy as jnp
from jax._src.api import value_and_grad
from jax import grad, jit, vmap, pmap, lax
# again, this only works on startup!
from jax.config import config
from jax.scipy.special import logsumexp
config.update('jax_platform_name', 'cpu')
config.update("jax_debug_nans", True)
config.update("jax_enable_x64", True)


@jit
def sigmoid(x):
    sigmoid_slope = 32.
    return 1.0 / (1.0 + jnp.exp(-sigmoid_slope * x))

@jit
def sigmoid2(x, x_offset):
    sigmoid_slope = 32.
    return 1. / (1. + jnp.exp(-sigmoid_slope * (x - x_offset)))

@jit
def in_range_unscaled(x, from_, to):
    return sigmoid2(x, from_)*(1. - sigmoid2(x, to))

@jit
def in_range(x, from_, to):
    return in_range_unscaled(x, from_, to) / in_range_unscaled((from_ + to) / 2.0, from_, to) # second term just for scaling

@jit
def sigmoid_with_slope(x, slope):
    return 1.0 / (1.0 + jnp.exp(-slope * x))

@jit
def periodic_step(x, period, offset, slope):
    return sigmoid_with_slope(jnp.sin((x - offset) * jnp.pi / period), slope)

def light_red(x, on_period, offset, slope = -1.0):
    if slope < 0.0:
      slope = 32.
    return periodic_step(x, on_period, offset, slope)

def min_pos_delta_if_equal(x, num_elems, ref_idx, y, z):
    sum = 0.0
    M = 1000.
    # def cond(pred, true_fun, operand):
    #     if pred:
    #         return body_fun(operand)
    # def body_fun(i, sum):
    #     if i != ref_idx:
    #         diff = x[i] - x[ref_idx]
    #         neg_offset = 2 * M * (1. - sigmoid(diff))
    #         diff += neg_offset
    #         lane_shift = 1000. * (1.0 - in_range(y[i] - z, -0.5, 0.5))
    #         diff += lane_shift
    #         exp_arg = -diff
    #         sum += jnp.exp(exp_arg)
    #         return sum
    # sum = lax.fori_loop(0, num_elems, body_fun, 0)
    # return -jnp.log(sum)

    for i in range(num_elems):
        if i == ref_idx: continue
        diff = x[i] - x[ref_idx]
        neg_offset = 2 * M * (1.0 - sigmoid(diff))
        diff += neg_offset
        lane_shift = 1000. * (1.0 - in_range(y[i] - z, -0.5, 0.5))
        diff += lane_shift
        exp_arg = -diff
        sum += jnp.exp(exp_arg)

    return -jnp.log(sum)
    # return -jnp.logsumexp((x[:num_elems]-x[ref_idx]) + 2*M*(1.0-sigmoid(x[i]-x[ref_idx]) + (1000. * (1. - in_range(y[:num_elems] - z, -.5, .5)))))


def max_pos(x, num_elems):
    return logsumexp(x[:num_elems])

def select_x_where_y_equals_z(x, y, z, num_elems):
    range_scale = 1.0
    return (x[:num_elems] * in_range(y[:num_elems] - z, -range_scale, range_scale)).sum()

def select_x_where_y0_equals_z0_and_y1_equals_z1(x, y0, z0, y1, z1, num_elems):
    return (x[:num_elems] * in_range(y0[:num_elems] - z0, -1.0, 1.0) * in_range(y1[:num_elems] - z1, -0.1, 0.1)).sum()

@jit
def select_x_if_larger_y_else_z(x, y, z):
    return sigmoid2(x, y) * x + (1.0 - sigmoid2(x, y)) * z

@jit
def idm(vel, leader_dist, leader_vel, speed_limit, max_accel):
    s0 = 5.0  # min spacing
    t = 3.0;   # desired time headway
    accel_term0 = jnp.power(vel / speed_limit, 4)
    accel_term1 = jnp.power((s0 + vel * t + (vel * (vel - leader_vel) / (2 * max_accel))) / leader_dist, 2)
    return max_accel * (1.0 - accel_term0 - accel_term1)

def traffic_sim_smoothed(x, num_agents_per_road):
    # Defining locals
    speed_limit = 13.888
    max_ts = 10.
    num_lanes = 3
    num_roads = 1
    tau = .1
    light_red_duration = 5.
    max_accel = 2.
    traffic_sim_sigmoid_slope = 16.0
    light_pos = 100.
    traffic_light_ahead_sigmoid_slope = 16.0
    # End definitions
    num_agents_per_road += num_lanes
    pos = jnp.zeros((num_roads, num_agents_per_road))
    vel = jnp.zeros((num_roads, num_agents_per_road))
    lane = jnp.zeros((num_roads, num_agents_per_road))

    sum_progress = 0.0
    light_offset = x

    step = 0
    while tau * step <= max_ts: 
        curr_light_red = light_red(tau * step, light_red_duration, light_offset, traffic_sim_sigmoid_slope)

        for road in range(num_roads):
            for i in range(num_agents_per_road - num_lanes):
                leader_dist = min_pos_delta_if_equal(pos[road], num_agents_per_road, i, lane[road], lane[road][i])
                leader_vel = select_x_where_y0_equals_z0_and_y1_equals_z1(vel[road], lane[road], lane[road][i], pos[road], pos[road][i] + leader_dist, num_agents_per_road)
                light_dist = light_pos - pos[road, i]
                light_ahead = sigmoid_with_slope(light_dist, traffic_light_ahead_sigmoid_slope)
                braking_at_light = curr_light_red * light_ahead * sigmoid_with_slope(leader_dist - light_dist, traffic_sim_sigmoid_slope)
                leader_dist = braking_at_light * light_dist + (1. - braking_at_light) * leader_dist
                accel = idm(vel[road, i], leader_dist, leader_vel, speed_limit, max_accel)
                vel = vel.at[road, i].add(accel.primal * tau)
                progress = vel[road, i] * tau
                sum_progress += progress
        step += 1
    return sum_progress

if __name__ == "__main__":
    num_agents_per_road = 1
    tau = .1
    x_init = tau/4 # Initial x time step
    print("x, y, dydx")
    value, gradient = value_and_grad(traffic_sim_smoothed)(x_init, num_agents_per_road)
    print(str(x_init) + ", ", value, ", " + str(gradient))

I have spent a good amount of time on this program in attempts for it to return a nonzero-grad value that isn't 0 when there should be a non-zero grad value when comparing it to the implementation in C++.

[mjhoover1]

Also, this is my attempt to vectorize the function, but I don't know how to handle the if i == ref_idx condition:

return -jnp.logsumexp( (x[:num_elems]-x[ref_idx]) + 
                                      (2*M*(1.0-sigmoid(x[i]-x[ref_idx]) + 
                                      (1000. * (1. - in_range(y[:num_elems] - z, -.5, .5)))
                                     )

[jakevdp]

For future reference, it's much better to put together a minimal reproducible example when asking questions like this. Your complete example mostly consists of code that is irrelevant to your question.

If you can provide a complete minimal example, I'd be happy to help! It looks like this function can be vectorized pretty straightforwardly.

[mjhoover1]

Okay, sorry about that @jakevdp I edited my question above to make it a minimal reproducible example, thanks for the advice!