C51 Projection Index Pair

Problem

Find the two support atoms a projected C51 target would land between.

Compute b = (value - v_min) / dz with dz = (v_max - v_min) / (n_atoms - 1), then return [lower_index, upper_index, upper_weight] where upper_weight = b - lower_index. Values at or below v_min give [0, 1, 0.0], values at or above v_max give [n_atoms - 1, n_atoms - 1, 0.0], and fewer than two atoms give [0, 0, 0.0].

Hint

The upper weight is the fractional part of the support coordinate b.

Test cases

case 1

in: [1.7,0,3,4]

expected: [1,2,0.7]

case 2

in: [-0.5,0,3,4]

expected: [0,1,0]

case 3

in: [3,0,3,4]

expected: [3,3,0]

case 4

in: [2.4,0,3,4]

expected: [2,3,0.3999999999999999]

case 5

in: [1,0,1,1]

expected: [0,0,0]

Starter codePython

def c51_projection_indices(value, v_min, v_max, n_atoms):
    # Your code here
    pass

5 test cases · run your code against them in the editor