Introduction
For a few years now, I’ve been working on setting a fast but general purpose way to recover imaging volumes from light field microscopy data, for getting data for fly muscle imaging experiments.
After some experimentation with a few light field toolboxes, we found that the oLaF toolbox worked really well for reconstructing volumes from our images.
Below you can see an example of an image from a light field microscope and a video of the reconstructed volume from deconvolution.
oLaF does work really well, but it is frustratingly slow. The reconstruction for the above image took 20 minutes. In our case, we want to process videos with thousands of frames, which would end up taking a week or more to process. To fix this, about two years ago, I spent 2 months porting the key deconvolution code in oLaF from Matlab to Python and then optimizing the code to run quickly on a GPU. This allowed us to process the above image in about 30 seconds and brings the video processing time down to just hours.
However, the code to calibrate the light field and calculate the forward and backward projection matrices for deconvolution was still in oLaF in Matlab. This wasn’t a big deal, this was part just needed to be run once and was relatively fast, but it was a bit annoying to have to Matlab for one part and then transfer to Python for the other. The calculations were in 3000 lines of heavy numerical Matlab code and I couldn’t justify the spare the time to translate this for our project.
Last week, I started playing around with ChatGPT to help me write and debug code. I played around with it and found that it could actually translate code from Matlab to Python really well, so translating the code in a few days suddenly seemed feasible. Why not try?
Workspace setup
I’ve been using the chat.el Emacs package to interact with ChatGPT. This made it really easy to send code over, since I could just select a region and press a few keys.
The translation begins!
I started from the main file which calculated camera and matrix parameters. The key code is:
%% Specific LFM configuration and camera parameters (um units)
Camera = LFM_setCameraParams(configFile, newSpacingPx);
%% Compute LFPSF Patterns and other prerequisites: lenslets centers, resolution related
[LensletCenters, Resolution, LensletGridModel, NewLensletGridModel] = ...
LFM_computeGeometryParameters(Camera, WhiteImage, depthRange, depthStep, superResFactor, 0);
[H, Ht] = LFM_computeLFMatrixOperators(Camera, Resolution, LensletCenters);As a test, I sent the code above with the prompt “translate the following matlab code to python” right above it. It sent me back the following Python code:
# Specific LFM configuration and camera parameters (um units)
Camera = LFM_setCameraParams(configFile, newSpacingPx)
# Compute LFPSF Patterns and other prerequisites: lenslets centers, resolution related
LensletCenters, Resolution, LensletGridModel, NewLensletGridModel = \
LFM_computeGeometryParameters(
Camera, WhiteImage, depthRange, depthStep, superResFactor, DebugBuildGridModel)
H, Ht = LFM_computeLFMatrixOperators(Camera, Resolution, LensletCenters)Okay that worked! But can it handle something more challenging? How about this snippet from the first function in the main file?
function Camera = LFM_setCameraParams(configFile, newSpacingPx)
Camera = ReadYaml(configFile);
%% compute extra LFM configuration specific parameters
% for regular grid we only need to compute 1/4 of the psf pattens, due to symmetry
if(strcmp(Camera.gridType, 'reg'))
Camera.range = 'quarter';
elseif(strcmp(Camera.gridType, 'hex'))
Camera.range = 'full';
end
Camera.fobj = Camera.ftl/Camera.M; %% focal length of objective lens
Camera.Delta_ot = Camera.ftl + Camera.fobj; %% obj2tube distance
% ulens spacing = ulens pitch
spacingPx = Camera.lensPitch/Camera.pixelPitch;
if(strcmp(newSpacingPx, 'default'))
newSpacingPx = spacingPx;
end
Camera.spacingPx = spacingPx;
Camera.newSpacingPx = newSpacingPx; % artificial spacing, usually lower than spacingPx for speed up
Camera.newPixelPitch = (Camera.lensPitch/newSpacingPx);After a bit of experimentation, I found that I had to prompt it with “translate the following matlab code to python, preserving comments” to get Python code with comments:
import yaml
def LFM_setCameraParams(configFile, newSpacingPx):
with open(configFile, 'r') as f:
Camera = yaml.safe_load(f)
## compute extra LFM configuration specific parameters
# for regular grid we only need to compute 1/4 of the psf pattens, due to symmetry
if Camera['gridType'] == 'reg':
Camera['range'] = 'quarter'
elif Camera['gridType'] == 'hex':
Camera['range'] = 'full'
Camera['fobj'] = Camera['ftl'] / Camera['M'] ## focal length of objective lens
Camera['Delta_ot'] = Camera['ftl'] + Camera['fobj'] ## obj2tube distance
# ulens spacing = ulens pitch
spacingPx = Camera['lensPitch'] / Camera['pixelPitch']
if newSpacingPx == 'default':
newSpacingPx = spacingPx
Camera['spacingPx'] = spacingPx
Camera['newSpacingPx'] = newSpacingPx # artificial spacing, usually lower than spacingPx for speed up
Camera['newPixelPitch'] = (Camera['lensPitch'] / newSpacingPx)
return CameraImpressive! Note that it translated the Matlab struct into a Python dictionary and used the pyyaml library to load the yaml file! Once the code gets to this level of complexity (which is honestly not too complex still), you start to get some randomness in the results. I ran it a few times, and once it didn’t properly translate the yaml loading (keeping the ReadYaml function as is) and another time it didn’t translate to the dictionary convention.
Getting more complex
A lot of the code in oLaF is of the kind above, which calls other functions, does relatively simple calculations, or organizes data. Generally GPT works well on this kind of code. How well does it work on code with more processing operations?
There it had some pretty mixed results. The biggest challenge I found was in translating the LFM_BuildLensletGridModel function. This is about 200 lines of Matlab code which calls several special Matlab functions with no direct Python numpy/scipy translation (fspecial to make a disk, DelaunayTri to query nearest neighbors efficiently, and imregionalmax to compute local maxes). I tried sending the whole function to the GPT a few times, but it would always make multiple mistakes in the translation for one reason or another.
Eventually, I found I could send it chunks of code about 40 lines long and it could translate these well enough, using the prompt “translate this to python, using skimage, scipy, and numpy libraries. make sure to think about indexing and translate .obj struct as [”obj”] access”. In the end, I still had to rewrite large chunks of this function manually to make it work well, taking some hints from the initial code on which functions to use.
Issues with speed
The rest of the code translated okay and I was able to run the whole pipeline within a day of starting, which was really exciting! However, problems quickly became apparent. The first issue was that the calculation of the forward projection matrix was extremely slow, about 20 times slower than Matlab! This made it hard to run more tests, so I tackled this first.
By looking through the code, I traced the slowness to this for loop in LFM_calcPSF:
for a in range(centerPT):
patternLine = zeroline.copy()
y = yspace[a]
for b in range(a, centerPT):
x = xspace[b]
xL2normsq = (((y + Camera['M'] * p1) ** 2 + (x + Camera['M'] * p2) ** 2)**0.5) / Camera['M']
# radial and axial optical coodinates
v = k * xL2normsq * np.sin(alpha)
# compute PSF;
# integrate over theta: 0->alpha
intgrand = lambda theta: np.sqrt(np.cos(theta)) * (1 + np.cos(theta)) * np.exp((1j * u/2) * (np.sin(theta/2) ** 2) / (np.sin(alpha/2) ** 2)) * jv(0, np.sin(theta) / np.sin(alpha) * v) * np.sin(theta) # zeroth order Bessel function of the first kind
I0, _ = quad(intgrand, 0, alpha)
patternLine[0, b] = Koi * I0
pattern[a, :] = patternLineI sent this to ChatGPT and asked just “how can i make this faster?”. It suggested to vectorize this loop, so I asked it for an implementation. It came out with this code, which really blew my mind:
y, x = np.meshgrid(yspace, xspace)
r = np.sqrt((y + Camera['M'] * p1) ** 2 + (x + Camera['M'] * p2) ** 2) / Camera['M']
v = k * r * np.sin(alpha)
theta = np.linspace(0, alpha, centerPT).reshape(-1, 1) # Make theta a column vector
intgrand = np.sqrt(np.cos(theta)) * (1 + np.cos(theta)) * np.exp((1j * u/2) * (np.sin(theta/2) ** 2) / (np.sin(alpha/2) ** 2)) * jv(0, np.sin(theta) / np.sin(alpha) * v) * np.sin(theta)
I0_re, _ = quad_vec(intgrand.real, 0, alpha) # quad_vec is a vectorized version of quad
I0_im, _ = quad_vec(intgrand.imag, 0, alpha)
Koi = demag / ((d1 * Camera['WaveLength']) ** 2) * np.exp(-1j * u / (4 * (np.sin(alpha/2) ** 2)))
# Vectorized multiplication with broadcasting
pattern = Koi * (I0_re + 1j * I0_im)It does have some issues, like quad_vec can actually handle complex integrands, intgrand should actually be a function, and pattern is not properly indexing alpha. Still, it’s so cool to see a double for loop (almost) automatically vectorized semantically. I was able to rewrite the code to this:
# create meshgrid of x and y values
x, y = np.meshgrid(xspace, yspace)
xL2normsq = (((y+Camera['M']*p1)**2+(x+Camera['M']*p2)**2)**0.5)/Camera['M']
# radial and axial optical coodinates
v = k*xL2normsq*np.sin(alpha)
def intgrand(theta, alpha, u, v):
return np.sqrt(np.cos(theta)) * (1+np.cos(theta)) * \
np.exp((1j*u/2) * (np.sin(theta/2)**2) / (np.sin(alpha/2)**2)) * \
scipy.special.j0(np.sin(theta)/np.sin(alpha)*v) * (np.sin(theta))
# integrate intgrand() over theta for all x and y values
# compute PSF
# integrate over theta: 0->alpha
I0x, _ = integrate.quad_vec(intgrand, 0, alpha,
args=(alpha, u, v), limit=100)
# compute pattern for all a and b values
pattern = np.zeros((centerPT, centerPT), dtype='complex128')
for a in range(centerPT):
pattern[a,a:] = Koi * I0x[a, a:]Vectorizing the for loop like this sped up the code about 4 times. I found also that using the specialized j0 to compute the zeroth order Bessel function of the first kind instead of jv further sped up by 5 times, bringing the speed in line with Matlab.
Off-by-one errors
So after all this, I was able to test the code and run it all the way through. However, when I did run it and compared the calculated projection matrices to the ones from oLaF, they were not the same! It turns out there were a lot of issues with off-by-one errors, where the original code expected 1-indexed variables but the python code expected 0-indexed variables. This might sound obvious, but in the context of this code which relied heavily on indexes and coordinate transformations, the bugs turned out to be frustratingly subtle. This took me the most time to debug (about 2 days, compared to 1 day for the rest above) and GPT was utterly useless in helping, as the bugs were based on subtle interactions across the whole codebase.
To demonstrate these kinds of interactions, here is an example with an off-by-one error that took me hours to properly find and debug. The following is a snippet from the LFM_computeForwardPatternsWaves in oLaF:
H = cell( coordsRange(1), coordsRange(2), length(Resolution.depths) );
for c = 1:length(Resolution.depths)
psfREF = psfWaveStack(:,:,c);
for aa_tex = 1:coordsRange(1)
aa_sensor = aa_tex / Resolution.texScaleFactor(1);
parfor bb_tex = 1:coordsRange(2)
bb_sensor = bb_tex / Resolution.texScaleFactor(2);
% shift the native plane PSF at every (aa, bb) position (native plane PSF is shift invariant)
psfSHIFT = imShift2(psfREF, round((aa_sensor-Nnum_half_coord(1))), round((bb_sensor-Nnum_half_coord(2))) );
% MLA transmittance
psfMLA = psfSHIFT.*MLARRAY;
% propagate the response to the sensor via Rayleigh-Sommerfeld diffraction
LFpsfAtSensor = prop2Sensor(psfMLA, sensorRes, Camera.mla2sensor, Camera.WaveLength, 0);
% shift the response back to center (we need the patterns centered for convolution)
LFpsf = imShift2(LFpsfAtSensor, round(-(aa_sensor-Nnum_half_coord(1))), round(-(bb_sensor-Nnum_half_coord(2))) );
% store the response pattern
H{aa_tex,bb_tex,c} = sparse(abs(double(LFpsf).^2));
end
end
endInitially, when I prompted GPT, it translated it to:
H = {}
for c in range(len(Resolution['depths'])):
psfREF = psfWaveStack[:,:,c]
for aa_tex in range(coordsRange[0]):
aa_sensor = aa_tex / Resolution['texScaleFactor'][0]
for bb_tex in range(coordsRange[1]):
bb_sensor = bb_tex / Resolution['texScaleFactor'][1]
# shift the native plane PSF at every (aa, bb) position (native plane PSF is shift invariant)
psfSHIFT = imShift2(psfREF, (aa_sensor-Nnum_half_coord[0]).round(), (bb_sensor-Nnum_half_coord[1]).round())
# MLA transmittance
psfMLA = psfSHIFT*MLARRAY
# propagate the response to the sensor via Rayleigh-Sommerfeld diffraction
LFpsfAtSensor = prop2Sensor(psfMLA, sensorRes, Camera['mla2sensor'], Camera['WaveLength'], False)
# shift the response back to center (we need the patterns centered for convolution)
LFpsf = imShift2(LFpsfAtSensor, (-(aa_sensor-Nnum_half_coord[0])).round(), (-(bb_sensor-Nnum_half_coord[1])).round())
# store the response pattern
H[(aa_tex,bb_tex,c)] = csr_matrix(np.abs(LFpsf**2))Do you see the error? It took me so long to see it!
It actually comes down to each of the aa_tex and bb_tex variables being used as both a coordinate and an index. As an index, they should be 0-indexed in Python and so the range in the for loops is correct. However, as a coordinate the matlab code expects 1-indexed coordinates. Generally it doesn’t matter much, but here it’s divided by Resolution['texScaleFactor'] so the off-by-one issue really affects it.
So the relevant section in the above code should really be:
for aa_tex in range(coordsRange[0]):
aa_sensor = (aa_tex + 1) / Resolution['texScaleFactor'][0]
for bb_tex in range(coordsRange[1]):
bb_sensor = (bb_tex + 1) / Resolution['texScaleFactor'][1]Adding 1 anywhere else would be wrong. You can’t add 1 to aa_sensor directly for instance, since the 1 is scaled by Resolution['texScaleFactor']. You also can’t only change the range(coordsRange[0]) to start at 1 either, since in H[(aa_tex,bb_tex,c)] = csr_matrix(np.abs(LFpsf**2)), the aa_tex is used as an index and should be 0-indexed as per Python convention.
There were about dozen bugs like this throughout the code. It took me a long time to really understand the root cause here. Once I recognized what it was though, it was relatively straightforward to fix it throughout the code.
Lessons
So what did I take away from all this? Overall, I think GPT worked pretty well to translate the code from Matlab to Python. It was really helpful in speeding up the tedious part of the translation and sometimes amazed me at helping with some more complex parts as well! Still, if you’ve read this far, you can see how it’s definitely not at the stage where you can just give it a code project and have it translate it from Matlab to Python flawlessly.
Here are some tricks that helped interact with this model:
For helping with code, GPT works best if you use it as a plugin from a text editor (Emacs in my case). It makes it easy to just keep sending it pieces of code and copy-paste back and forth. I also just got better code results from the plugin than from the web interface. I’m not sure why that is, it might have to do with the initial prompt of “You are a helpful AI Assistant embedded into Emacs.”
For interpreting/translating code, it starts breaking after ~100 lines. To solve this, it helps to feed it smaller chunks of code. Often times, functions are less than 100 lines so function-by-function translation works well. If a function is over 100 lines, it works to just split the function roughly semantically and feed it each of the chunks from this split.
Sometimes when it generates something, I spot an error right there. In that case, it often helps to follow with something like “there is an indexing error on the line below, please fix it: [line in question]”.
When I suspect there may be an error but I don’t directly spot it, I usually reply to the generate code response with “please review the translation above, point out errors, and fix the code. think step by step.” to make it review errors. It often corrects errors if they are present. If there’s no error, the model does go along with the prompt anyway and just tries to fix something but makes it worse. Be careful with this.
This was a really good use case for GPT and got me really familiar with using it for coding. I have a much better sense of how it helps and where it falls short now. I will likely keep it1 as another tool within my editor.
Conclusion
With all of this done, I was able to finish the translation of oLaF into a self-contained Python package with GPU acceleration: pyolaf. If you’re doing light field microscopy, try it out!