jax conversion of trace.py in ray-optics

[PoulomiPradhan]

I am trying to convert trace.py into jax. So for that instead of using numpy libraries, I am importing equivalent jax libraries. The original trace.py was using "Newton" and "fsolove" from scipy.optimize but jax.scipy doesn't have newton/fsolve. So I am trying to use "minimize" for the same operation. I have attached the code below, that's the only function I have changed.

In the sub function "y_stop_coordinate" ," pt1 = np.array([0., y1, dist])" This line is throwing ValueError "All input arrays must have the same shape" y1 is of type class <jax.interpreters.ad.JVPTracer> and dist is of type float.
out = stack([asarray(elt,dtype=dtype)]) in lax_numpy.py (line number 1900) ) is causing this.
Please suggest how can I overcome this error ? (Given that I want the files to be converted in jax for auto differentiation calculation)

import itertools
import logging

import pdb
import jax.numpy as np
from jax.numpy.linalg import norm
from jax.scipy.optimize import minimize

from jax import grad, jit, vmap
from jax.numpy import sqrt, copysign, sin
from jax.lax import acos

#from scipy.optimize import newton, fsolve
import pandas as pd

from . import raytrace as rt
from . import RayPkg, RaySeg
from .waveabr import (wave_abr_full_calc, calculate_reference_sphere,
transfer_to_exit_pupil)
from rayoptics.optical import model_constants as mc
from .traceerror import TraceError, TraceMissedSurfaceError, TraceTIRError

def iterate_ray(opt_model, ifcx, xy_target, fld, wvl, **kwargs):
    """ iterates a ray to xy_target on interface ifcx, returns aim points on
    the paraxial entrance pupil plane

    If idcx is None, i.e. a floating stop surface, returns xy_target.

    If the iteration fails, a TraceError will be raised
    """
    def y_stop_coordinate(y1, *args):
        seq_model, ifcx, pt0, dist, wvl, y_target = args
        
        #import pdb
        #pdb.set_trace()

        pt1 = np.array([0., y1, dist])

        dir0 = pt1 - pt0
        length = norm(dir0)
        dir0 = dir0/length
        try:
            ray, _, _ = rt.trace(seq_model, pt0, dir0, wvl)
        except TraceMissedSurfaceError as ray_miss:
            ray = ray_miss.ray_pkg
            if ray_miss.surf <= ifcx:
                raise ray_miss
        except TraceTIRError as ray_tir:
            ray = ray_tir.ray_pkg
            if ray_tir.surf < ifcx:
                raise ray_tir
        y_ray = ray[ifcx][mc.p][1]
    
        
        return y_ray - y_target

    def surface_coordinate(coord, *args):
        seq_model, ifcx, pt0, dist, wvl, target = args
        pt1 = np.array([coord[0], coord[1], dist])
        dir0 = pt1 - pt0
        length = norm(dir0)
        dir0 = dir0/length
        ray, _, _ = rt.trace(seq_model, pt0, dir0, wvl)
        xy_ray = np.array([ray[ifcx][mc.p][0], ray[ifcx][mc.p][1]])

    print(coord[0], coord[1], xy_ray[0], xy_ray[1])

        return xy_ray - target

    seq_model = opt_model.seq_model
    osp = opt_model.optical_spec

    fod = opt_model['analysis_results']['parax_data'].fod
    dist = fod.obj_dist + fod.enp_dist

    pt0 = osp.obj_coords(fld)
    if ifcx is not None:
        if pt0[0] == 0.0 and xy_target[0] == 0.0:
            # do 1D iteration if field and target points are zero in x
            y_target = xy_target[1]
            logging.captureWarnings(True)
            try:
                start_y, results = minimize(y_stop_coordinate,np.array([0.]),
                                        args=(seq_model, ifcx, pt0,
                                                dist, wvl, y_target),method = "BFGS"
                )

                                            
            #try:
            #   start_y, results = newton(y_stop_coordinate, 0.,
            #                            args=(seq_model, ifcx, pt0,
            #                                 dist, wvl, y_target),
                #                          disp=False, full_output=True)                              
            except RuntimeError as rte:
                # if we come here, start_y is a RuntimeResults object
                # print(rte)
                start_y = results.root
            except TraceError:
                start_y = 0.0
            start_coords = np.array([0., start_y])
        else:
            # do 2D iteration. epsfcn is a parameter increment,
            #  make proportional to pupil radius
            try:
                start_coords = minimize(surface_coordinate, np.array([0., 0.]),
                                    epsfcn=0.0001*fod.enp_radius,
                                    args=(seq_model, ifcx, pt0, dist,
                                            wvl, xy_target),method = "BFGS")
            except TraceError:
                start_coords = np.array([0., 0.])
    else:  # floating stop surface - use entrance pupil for aiming
        start_coords = np.array([0., 0.]) + xy_target

    return start_coords

[mjhoptics]

Hello,
The original code was passing a scalar (float) into y_stop_coordinate. You’re now passing a 1-D array of length 1 into it. In y_stop_coordinate, y1 is assumed to be a scalar and thus the exception. So you either need to pass a scalar into y_stop_coordinate or dereference the array inside of that function. I don’t know anything about Jax, other than what I just read in the doc, so I don’t know which fix is best/possible.

[PoulomiPradhan]

As you are saying " You’re now passing a 1-D array of length 1 into it.", do you point out this following line ?

            try:
                start_y, results = minimize(y_stop_coordinate,np.array([0.]),
                                        args=(seq_model, ifcx, pt0,
                                                dist, wvl, y_target),method = "BFGS"
                )

Because I haven't changed " pt1 = np.array([0., y1, dist])" in y_stop_coordinate function.

[mjhoptics]

I'm guessing the difference is that minimize() works on a vector, in this case of length 1, whereas newton() works on a scalar function. So you'll have to change y_stop_coordinate() to assume a vector is coming in, not a scalar.

Incidentally, when I originally implemented iterate_ray() I found that if all I needed was a 1-D iteration, which is a dominant use-case overall, that using different algorithms for the 1-D and 2-D cases made a lot of sense and was faster and more reliable. Since jax doesn't have a 1-D minimizer, you might just fall back to treating everything as a 2-D case.

[PoulomiPradhan]

Yes. I am passing np.array([0.]) because minimize() works on a vector. using pdb I checked the type pf y1 for my program is <jax.interpreters.ad.JVPTracer>and type of dist in regular float. The detail error is showing below :

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In [6], line 17
     13 d_L = d_L1 + d_L2 # lens thickness
     15 max_angle = 1/180*np.pi # max incident angle, rad
---> 17 sm = lens(f, R1, R2, d_L, wv, n, pupil_r)
     18 visualize(sm)
     19 plt.title('Without correction')

Cell In [4], line 28, in lens(f, R1, R2, d_L, wv, n, pupil_r, p)
     25 sm.add_surface([R2, 0])
     27 sm.gaps[-1].thi = remaining_height - d_L
---> 28 opm.update_model()
     29 return sm

File ~/.local/lib/python3.10/site-packages/rayoptics/optical/opticalmodel.py:333, in OpticalModel.update_model(self, **kwargs)
    331 self['seq_model'].update_model(**kwargs)
    332 self['optical_spec'].update_model(**kwargs)
--> 333 self.update_optical_properties(**kwargs)
    334 if len(self['part_tree'].nodes_with_tag(tag='#element')) == 0:
    335     elements_from_sequence(self['ele_model'], self['seq_model'], 
    336                            self['part_tree'])

File ~/.local/lib/python3.10/site-packages/rayoptics/optical/opticalmodel.py:348, in OpticalModel.update_optical_properties(self, **kwargs)
    345 """Compute first order and other optical properties. """
    347 # OpticalSpec maintains first order and ray aiming for fields
--> 348 self['optical_spec'].update_optical_properties(**kwargs)
    349 # Update the ParaxialModel as needed
    350 self['parax_model'].update_model(**kwargs)

File ~/.local/lib/python3.10/site-packages/rayoptics/raytr/opticalspec.py:169, in OpticalSpecs.update_optical_properties(self, **kwargs)
    167 if self.do_aiming:
    168     for i, fld in enumerate(self.field_of_view.fields):
--> 169         aim_pt = aim_chief_ray(self.opt_model, fld, wvl)
    170         fld.aim_pt = aim_pt

File ~/.local/lib/python3.10/site-packages/rayoptics/raytr/trace.py:527, in aim_chief_ray(opt_model, fld, wvl)
    525     wvl = seq_model.central_wavelength()
    526 stop = seq_model.stop_surface
--> 527 aim_pt = iterate_ray(opt_model, stop, np.array([0., 0.]), fld, wvl)
    528 return aim_pt

File ~/.local/lib/python3.10/site-packages/rayoptics/raytr/trace.py:304, in iterate_ray(opt_model, ifcx, xy_target, fld, wvl, **kwargs)
    302 logging.captureWarnings(True)
    303 try:
--> 304     start_y, results = minimize(y_stop_coordinate,np.array([0.]),
    305                               args=(seq_model, ifcx, pt0,
    306                                     dist, wvl, y_target),method = "BFGS"
    307     )
    310 #try:
    311  #   start_y, results = newton(y_stop_coordinate, 0.,
    312   #                            args=(seq_model, ifcx, pt0,
    313    #                                 dist, wvl, y_target),
    314     #                          disp=False, full_output=True)                              
    315 except RuntimeError as rte:
    316     # if we come here, start_y is a RuntimeResults object
    317     # print(rte)

File ~/.local/lib/python3.10/site-packages/jax/_src/scipy/optimize/minimize.py:103, in minimize(fun, x0, args, method, tol, options)
    100 fun_with_args = lambda x: fun(x, *args)
    102 if method.lower() == 'bfgs':
--> 103   results = minimize_bfgs(fun_with_args, x0, **options)
    104   success = results.converged & jnp.logical_not(results.failed)
    105   return OptimizeResults(x=results.x_k,
    106                          success=success,
    107                          status=results.status,
   (...)
    112                          njev=results.ngev,
    113                          nit=results.k)

File ~/.local/lib/python3.10/site-packages/jax/_src/scipy/optimize/bfgs.py:100, in minimize_bfgs(fun, x0, maxiter, norm, gtol, line_search_maxiter)
     97 d = x0.shape[0]
     99 initial_H = jnp.eye(d, dtype=x0.dtype)
--> 100 f_0, g_0 = jax.value_and_grad(fun)(x0)
    101 state = _BFGSResults(
    102     converged=jnp.linalg.norm(g_0, ord=norm) < gtol,
    103     failed=False,
   (...)
    114     line_search_status=0,
    115 )
    117 def cond_fun(state):

    [... skipping hidden 8 frame]

File ~/.local/lib/python3.10/site-packages/jax/_src/scipy/optimize/minimize.py:100, in minimize.<locals>.<lambda>(x)
     97   msg = "args argument to jax.scipy.optimize.minimize must be a tuple, got {}"
     98   raise TypeError(msg.format(args))
--> 100 fun_with_args = lambda x: fun(x, *args)
    102 if method.lower() == 'bfgs':
    103   results = minimize_bfgs(fun_with_args, x0, **options)

File ~/.local/lib/python3.10/site-packages/rayoptics/raytr/trace.py:260, in iterate_ray.<locals>.y_stop_coordinate(y1, *args)
    256 y1
    257 #import pdb
    258 #pdb.set_trace()
--> 260 pt1 = np.array([0., y1, dist])
    262 dir0 = pt1 - pt0
    263 length = norm(dir0)

File ~/.local/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py:1900, in array(object, dtype, copy, order, ndmin)
   1898 elif isinstance(object, (list, tuple)):
   1899   if object:
-> 1900     out = stack([asarray(elt, dtype=dtype) for elt in object])
   1901   else:
   1902     out = np.array([], dtype=dtype)

File ~/.local/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py:1640, in stack(arrays, axis, out, dtype)
   1637 for a in arrays:
   1638   if shape(a) != shape0:
-> 1640     raise ValueError("All input arrays must have the same shape.")
   1641   new_arrays.append(expand_dims(a, axis))
   1642 return concatenate(new_arrays, axis=axis, dtype=dtype)

ValueError: All input arrays must have the same shape.

So what should I change ?

[mjhoptics]

Try changing

pt1 = np.array([0., y1, dist])

to

pt1 = np.array([0., y1[0], dist])

If that doesn't work, I'm not clear what the problem is.

[PoulomiPradhan]

Thanks for your answer, but my error just changed after I changed to pt1 = np.array([0., y1[0], dist])

---------------------------------------------------------------------------
TracerArrayConversionError                Traceback (most recent call last)
Cell In [6], line 17
     13 d_L = d_L1 + d_L2 # lens thickness
     15 max_angle = 1/180*np.pi # max incident angle, rad
---> 17 sm = lens(f, R1, R2, d_L, wv, n, pupil_r)
     18 visualize(sm)
     19 plt.title('Without correction')

Cell In [4], line 28, in lens(f, R1, R2, d_L, wv, n, pupil_r, p)
     25 sm.add_surface([R2, 0])
     27 sm.gaps[-1].thi = remaining_height - d_L
---> 28 opm.update_model()
     29 return sm

File ~/.local/lib/python3.10/site-packages/rayoptics/optical/opticalmodel.py:333, in OpticalModel.update_model(self, **kwargs)
    331 self['seq_model'].update_model(**kwargs)
    332 self['optical_spec'].update_model(**kwargs)
--> 333 self.update_optical_properties(**kwargs)
    334 if len(self['part_tree'].nodes_with_tag(tag='#element')) == 0:
    335     elements_from_sequence(self['ele_model'], self['seq_model'], 
    336                            self['part_tree'])

File ~/.local/lib/python3.10/site-packages/rayoptics/optical/opticalmodel.py:348, in OpticalModel.update_optical_properties(self, **kwargs)
    345 """Compute first order and other optical properties. """
    347 # OpticalSpec maintains first order and ray aiming for fields
--> 348 self['optical_spec'].update_optical_properties(**kwargs)
    349 # Update the ParaxialModel as needed
    350 self['parax_model'].update_model(**kwargs)

File ~/.local/lib/python3.10/site-packages/rayoptics/raytr/opticalspec.py:169, in OpticalSpecs.update_optical_properties(self, **kwargs)
    167 if self.do_aiming:
    168     for i, fld in enumerate(self.field_of_view.fields):
--> 169         aim_pt = aim_chief_ray(self.opt_model, fld, wvl)
    170         fld.aim_pt = aim_pt

File ~/.local/lib/python3.10/site-packages/rayoptics/raytr/trace.py:527, in aim_chief_ray(opt_model, fld, wvl)
    525     wvl = seq_model.central_wavelength()
    526 stop = seq_model.stop_surface
--> 527 aim_pt = iterate_ray(opt_model, stop, np.array([0., 0.]), fld, wvl)
    528 return aim_pt

File ~/.local/lib/python3.10/site-packages/rayoptics/raytr/trace.py:304, in iterate_ray(opt_model, ifcx, xy_target, fld, wvl, **kwargs)
    302 logging.captureWarnings(True)
    303 try:
--> 304     start_y, results = minimize(y_stop_coordinate,np.array([0.]),
    305                               args=(seq_model, ifcx, pt0,
    306                                     dist, wvl, y_target),method = "BFGS"
    307     )
    310 #try:
    311  #   start_y, results = newton(y_stop_coordinate, 0.,
    312   #                            args=(seq_model, ifcx, pt0,
    313    #                                 dist, wvl, y_target),
    314     #                          disp=False, full_output=True)                              
    315 except RuntimeError as rte:
    316     # if we come here, start_y is a RuntimeResults object
    317     # print(rte)

File ~/.local/lib/python3.10/site-packages/jax/_src/scipy/optimize/minimize.py:103, in minimize(fun, x0, args, method, tol, options)
    100 fun_with_args = lambda x: fun(x, *args)
    102 if method.lower() == 'bfgs':
--> 103   results = minimize_bfgs(fun_with_args, x0, **options)
    104   success = results.converged & jnp.logical_not(results.failed)
    105   return OptimizeResults(x=results.x_k,
    106                          success=success,
    107                          status=results.status,
   (...)
    112                          njev=results.ngev,
    113                          nit=results.k)

File ~/.local/lib/python3.10/site-packages/jax/_src/scipy/optimize/bfgs.py:100, in minimize_bfgs(fun, x0, maxiter, norm, gtol, line_search_maxiter)
     97 d = x0.shape[0]
     99 initial_H = jnp.eye(d, dtype=x0.dtype)
--> 100 f_0, g_0 = jax.value_and_grad(fun)(x0)
    101 state = _BFGSResults(
    102     converged=jnp.linalg.norm(g_0, ord=norm) < gtol,
    103     failed=False,
   (...)
    114     line_search_status=0,
    115 )
    117 def cond_fun(state):

    [... skipping hidden 8 frame]

File ~/.local/lib/python3.10/site-packages/jax/_src/scipy/optimize/minimize.py:100, in minimize.<locals>.<lambda>(x)
     97   msg = "args argument to jax.scipy.optimize.minimize must be a tuple, got {}"
     98   raise TypeError(msg.format(args))
--> 100 fun_with_args = lambda x: fun(x, *args)
    102 if method.lower() == 'bfgs':
    103   results = minimize_bfgs(fun_with_args, x0, **options)

File ~/.local/lib/python3.10/site-packages/rayoptics/raytr/trace.py:266, in iterate_ray.<locals>.y_stop_coordinate(y1, *args)
    264 dir0 = dir0/length
    265 try:
--> 266     ray, _, _ = rt.trace(seq_model, pt0, dir0, wvl)
    267 except TraceMissedSurfaceError as ray_miss:
    268     ray = ray_miss.ray_pkg

File ~/.local/lib/python3.10/site-packages/rayoptics/raytr/raytrace.py:81, in trace(seq_model, pt0, dir0, wvl, **kwargs)
     78 kwargs['first_surf'] = kwargs.get('first_surf', 1)
     79 kwargs['last_surf'] = kwargs.get('last_surf',
     80                                  seq_model.get_num_surfaces()-2)
---> 81 return trace_raw(path, pt0, dir0, wvl, **kwargs)

File ~/.local/lib/python3.10/site-packages/rayoptics/raytr/raytrace.py:161, in trace_raw(path, pt0, dir0, wvl, eps, check_apertures, **kwargs)
    158 surf += 1
    160 rt, t = tfrm_from_before
--> 161 b4_pt, b4_dir = rt.dot(before_pt - t), rt.dot(before_dir)
    163 pp_dst = -b4_pt.dot(b4_dir)
    164 pp_pt_before = b4_pt + pp_dst*b4_dir

File ~/.local/lib/python3.10/site-packages/jax/core.py:540, in Tracer.__array__(self, *args, **kw)
    539 def __array__(self, *args, **kw):
--> 540   raise TracerArrayConversionError(self)

TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ConcreteArray([ 0.  0. -1.], dtype=float32)>with<JVPTrace(level=2/0)> with
  primal = DeviceArray([ 0.,  0., -1.], dtype=float32)
  tangent = Traced<ShapedArray(float32[3])>with<JaxprTrace(level=1/0)> with
    pval = (ShapedArray(float32[3]), None)
    recipe = JaxprEqnRecipe(eqn_id=<object object at 0x7f57e23caab0>, in_tracers=(Traced<ShapedArray(float32[3]):JaxprTrace(level=1/0)>,), out_tracer_refs=[<weakref at 0x7f57da0e6070; to 'JaxprTracer' at 0x7f57da0e65c0>], out_avals=[ShapedArray(float32[3])], primitive=xla_call, params={'device': None, 'backend': None, 'name': '<lambda>', 'donated_invars': (False,), 'inline': True, 'keep_unused': False, 'call_jaxpr': { lambda ; a:f32[3]. let  in (a,) }}, effects=set(), source_info=SourceInfo(traceback=<jaxlib.xla_extension.Traceback object at 0x7f57da0efcb0>, name_stack=NameStack(stack=(Transform(name='jvp'),))))
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError

[mjhoptics]

It appears you are running into type conversion issues between numpy and jax. The error message contains a link to the jax documentation concerning this. I would start there and investigate which part of the expression is causing the error and then seeing if changing the types of the expression pieces fixes the problem.

I don't know anything about jax but replacing a calculation engine (numpy) with another engine (jax) in central ray trace code is not a simple task, I suspect. I would trace as simple a model as possible so that you can get intermediate values from the numpy code to compare with the jax code.

Good luck!

[PoulomiPradhan]

Thanks. Yes. Now I am trying to change all the Ray-optics files into jax.numpy instead of numpy. After changing some files, the error changed into ConcretizationTypeError(tracer, context=''). I will post the exact error after changing some other files into jax.numpy.

[mjhoptics]

Replacing the computational engine that ray-optics uses might best be approached from the bottom-up, in terms of the program structure. For ray-optics, the core ray trace calculations are done in the raytrace module and the profiles module. The trace_raw function is the core ray trace loop. There are functions for computing the optical interaction (bend, reflect, phase) in the raytrace module, while the methods for intersecting a ray with a surface are implemented in the profiles module.

The sequential path definition is the key data structure that the ray trace loop runs over. The sequential path is a list of tuples, each tuple consisting of (using python type hints):

Path_Seg = Tuple[ifc: 'interface.Interface', gp: 'gap.Gap', lcl_tfrm: Transform3, rndx: float, z_dir: ZDir]
Path = Iterator[Path_Seg]

where:

Vector3 = np.ndarray
Matrix3 = np.ndarray
Transform3 = Tuple[Matrix3, Vector3]
ZDir = Literal[1, -1]

The path and reverse_path methods of the sequential model return Path iterators that return Path_Segs. The main loop in trace_raw loops over Path_Segs.

The lcl_tfrm data is the transformation needed to go from the interface i coordinate system to the i+1 coordinates. It consists of a 3D rotation matrix and a 3D translation vector. It is calculated by forward_transform, based on data in ifc.decenter_data. The lcl_tfrm is calculated from DecenterData. I mention this in particular since the error reported above occurs when processing the lcl_tfrm data (called tfrm_from_before in trace_raw).

An example of a "minimal" code ray trace test is in the rayoptics/raytr/tests directory called test_sequential.py. This uses a function gen_sequence to provide the sequential path definition for the trace_raw call. This example touches the smallest footprint of the ray-optics code while traversing most paths through the core ray trace code.

Once the low level ray trace is working, the next area to include might be the analyses module, for tracing grids and lists of rays.

I think what you're trying to do is interesting and ambitious. For a code base like ray-optics, it is important to understand how the modules fit together and how they are layered, from computation outward to the user interface. If I can help you by documenting how these different things fit together (and which things you don't have to worry about), I'm more than happy to answer questions and fill in details. But I can't do line-by-line debugging; I hope you understand.
Mike Hayford