Laser physics simulation

The speed is quite slow, which seems to be delayed a lot by the three nested for loops with indices j,k, and l

This is a good start to finding out which part of the code is causing a bottleneck. Guessing which part is slow can be good for building up an intuition, so let's test the code and see what the slowdown is.

We can measure the performance of the code in a few ways. Since I have no understanding of the code, I would begin by throwing a general profiler at it. Python comes with a profiler called cProfile. Another profiler I find useful is pyinstrument.

Running the code through cProfile python -m cProfile -s time cr_laser_physics.py and killing it after a minute gave me a profile. The top of it looks like this

801051 function calls (786162 primitive calls) in 62.441 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1   45.036   45.036   62.804   62.804 cr_laser_physics.py:1(<module>)
13081/8182   13.425    0.001   15.052    0.002 {built-in method numpy.core._multiarray_umath.implement_array_function}
    26476    1.239    0.000    1.612    0.000 {built-in method builtins.eval}
    12644    1.195    0.000    1.195    0.000 {method 'reduce' of 'numpy.ufunc' objects}
     4719    0.226    0.000   13.868    0.003 numeric.py:943(tensordot)
  104/102    0.201    0.002    0.203    0.002 {built-in method _imp.create_dynamic}
      448    0.073    0.000    0.073    0.000 {method 'read' of '_io.FileIO' objects}
      448    0.072    0.000    0.072    0.000 {built-in method marshal.loads}
    14217    0.057    0.000    0.057    0.000 {method 'reshape' of 'numpy.ndarray' objects}
        1    0.034    0.034    0.034    0.034 {built-in method mkl._py_mkl_service.get_version}
...

Unfortunately, this doesn't look all that helpful. Since the majority of the time is spent in the module (45 seconds of it) and the remainder is a numpy method that we aren't calling directly, we haven't learned much. We can definitely say that the code overall is slow.

A google search lead me to this question. The problem they are having looks similar to ours, cProfile isn't being specific about where the performance problem is. The suggested answer is to split the code into more functions, so that cProfile can say which ones are slow.

So this looks like a good time to go through the code, and make whatever improvements we can. Some nice things to do would be

1. Remove unused imports. They add a small bit of noise to the profiler. Pylance suggests that most of the imports are not needed. You can always add them back later when needed. The special import is falsely marked as unused, because Pylance cannot tell that the strings of code will be run here. This is a good indicator that the code is doing something in a very strange way. Can this be simplified?
2. Put things into functions. While every function call will have an overhead, we need it so that cProfile can profile well.
3. Split long lines into smaller chunks. This will make it easier to see what parts of the code depend on what.

I won't implement all of these suggestions, but I'll try and outline what the file contents will look like.

# Imports
import copy
import math

import numpy as np
import scipy.special as special

from scipy.integrate import quad

# Global constants. (The usual convention is to make these all uppercase).
nz = 70
ny = nx = 3
ntau = 70
...

def main():
    S=np.empty((nx,ny,nz,nx,ny,nz,ntau),dtype='complex_')
    S[:,:,:,:,:,:,0]=0
    rho=np.empty((nx,ny,nz,ntau),dtype='complex_')
    ...


if __name__ == "__main__":
    main()

From this point, I would pick at the contents of the main method until it is split into as many functions as make sense. The hard part here is tracking which variables depend on which other variables.

Doing so, I end up with a main method that looks something like

def main():
    rho = make_rho()
    g = make_g()
    g2 = make_g2(g)

    Intensitytrans = big_loop(g, g2, rho)
    Intensity = make_intensity(Intensitytrans)

    Intensityn, photonn = smaller_loop(Intensity)

Here I've kept the variable the same, but moved as much code out into functions as I could. Each function follows a similar pattern of 'create variables -> loop to fill them with values -> return the variables referenced later'. As an example, here is make_g2.

def make_g2(g):
    g2=np.zeros((2*(nx-1)+1,2*(ny-1)+1,2*(nz-1)+1),dtype ='complex_')
    for i in range(2*(nz-1)+1):
        g2[:,:,i]=copy.copy(g[:,:,i])*(0 if i>nz else 1)
    return g2

Running the profiler on the code now gives: (cut to just show the 5 most expensive functions)

516776 function calls (504160 primitive calls) in 60.096 seconds

Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1   44.271   44.271   59.010   59.010 cr_laser_physics.py:75(big_loop)
19107/11944   12.793    0.001   14.640    0.001 {built-in method numpy.core._multiarray_umath.implement_array_function}
    19102    1.385    0.000    1.385    0.000 {method 'reduce' of 'numpy.ufunc' objects}
    26460    0.803    0.000    1.060    0.000 {built-in method builtins.eval}
     7163    0.236    0.000   13.265    0.002 numeric.py:943(tensordot)

This tells us two pieces of information. You were right that most of the time (45 seconds when the code is run for 60 seconds). is spent in the heavily-nested loops. The other functions aren't noticably slow when compared to those loops, so we don't need to worry about them now.

The slow function looks like this

@profile
def big_loop(g, g2, rho):
    # This has a lot going on, can it be simplified?
    S=np.empty((nx,ny,nz,nx,ny,nz,ntau),dtype='complex_')
    S[:,:,:,:,:,:,0]=0

    Intensitytrans=np.zeros((nx,ny,nz,ntau),dtype='complex_')

    for i in range(ntau):
        i1=np.zeros((nx,ny,nz),dtype='complex_')
        i2=np.zeros((nx,ny,nz),dtype='complex_')
        if i<ntau-1:
            snoise1=np.zeros((nx,ny,nz,nx,ny,nz),dtype='complex_')
            snoise2=np.zeros((nx,ny,nz,nx,ny,nz),dtype='complex_')
            snoise=np.zeros((nx,ny,nz,nx,ny,nz),dtype='complex_')
            s1=np.zeros((nx,ny,nz),dtype='complex_')
            s2=np.zeros((nx,ny,nz,nx,ny,nz),dtype='complex_')
            s3=np.zeros((nx,ny,nz,nx,ny,nz),dtype='complex_')
        for j in range(nx):
            for k in range(ny):
                for l in range(nz):  # THESE LOOPS MAY SLOW THE COMPUTATION
                    gc=copy.copy(g[(nx-1)-j:2*(nx-1)+1-j,(ny-1)-k:2*(ny-1)+1-k,(nz-1)-l:2*(nz-1)+1-l])
                    g2c=copy.copy(g2[(nx-1)-j:2*(nx-1)+1-j,(ny-1)-k:2*(ny-1)+1-k,(nz-1)-l:2*(nz-1)+1-l])
                    i1[j,k,l]=np.sum(S[:,:,:,:,:,:,i]*np.tensordot(np.conjugate(g2c),g2c,axes=0))*(dx*dy*dz)**2
                    i2[j,k,l]=np.sum((1+rho[:,:,:,i])/2*np.absolute(g2c)**2)*dx*dy*dz
                    if i<ntau-1:
                        s1+=S[:,:,:,j,k,l,i]*np.conjugate(g2c)*dx*dy*dz
                        s2+=np.tensordot(np.conjugate(g2c)*rho[:,:,:,i],S[j,k,l,:,:,:,i],axes = 0)*dx*dy*dz
                        s3+=np.tensordot(S[:,:,:,j,k,l,i],g2c*rho[:,:,:,i],axes = 0)*dx*dy*dz
                        snoise1[:,:,:,j,k,l]=rho[:,:,:,i]*(1+rho[j,k,l,i])/2*np.conjugate(gc)
                        snoise2[j,k,l,:,:,:]=rho[j,k,l,i]*(1+rho[:,:,:,i])/2*np.conjugate(gc)
        Intensitytrans[:,:,:,i]=3/(32*math.pi*lambda0**2*delo)*gamsp*(nnu**2*i1+nnu*i2)
        if i<ntau-1:
            for j in range(nz):
                for k in range(nz):
                    snoise[:,:,j,:,:,k]=snoise1[:,:,j,:,:,k] if j>k else snoise2[:,:,j,:,:,k]
            rho[:,:,:,i+1]=rho[:,:,:,i]+gamsp*(-2*(1+rho[:,:,:,i])/2-3*nnu/(8*math.pi)*2*np.real(s1))*dtau
            S[:,:,:,:,:,:,i+1]=S[:,:,:,:,:,:,i]+gamsp*(-S[:,:,:,:,:,:,i]+3/(16*math.pi)*(nnu*(s2+s3)+snoise))*dtau

    return Intensitytrans

Since splitting this function up doesn't look easy, the next thing to do is use a different profiler. A line-by-line profiler should be more helpful here.

I installed line_profiler through pip, added @profile to specify the function to profile, and ran it using kernprof -l cr_laser_physics.py, killing it after 60 seconds. I then viewed the profile using python -m line_profiler cr_laser_physics.py.lprof. Keeping only the lines which contribute more than 5% of the time shows three lines are really slow.

Timer unit: 1e-06 s

Total time: 60.4413 s
File: cr_laser_physics.py
Function: big_loop at line 74

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    74                                           @profile
    75                                           def big_loop(g, g2, rho):
    97      2433   22926489.0   9423.1     37.9                      i1[j,k,l]=np.sum(S[:,:,:,:,:,:,i]*np.tensordot(np.conjugate(g2c),g2c,axes=0))*(dx*dy*dz)**2
   101      2433   18182798.0   7473.4     30.1                          s2+=np.tensordot(np.conjugate(g2c)*rho[:,:,:,i],S[j,k,l,:,:,:,i],axes = 0)*dx*dy*dz
   102      2433   18384856.0   7556.5     30.4                          s3+=np.tensordot(S[:,:,:,j,k,l,i],g2c*rho[:,:,:,i],axes = 0)*dx*dy*dz

Each of those lines are 'hot', they take a long time and are run many many times. I'm not too confident about what can be changed, so I'll leave some suggestions below.

i1[j,k,l] = np.sum(S[:,:,:,:,:,:,i]*np.tensordot(np.conjugate(g2c),g2c,axes=0))*(dx*dy*dz)**2

Both np.sum(S[:,:,:,:,:,:,i] and (dx*dy*dz)**2 will always give the same result, even when j, k, and l change. They can be done once before that triple-nested loop. I suspect that the sum is very costly to compute so often.

Each of the slow lines has a call np.conjugate(g2c). I would recommend doing this once and storing it in a variable.

The latter two slow lines are both adding a tensor product to a sum. Could you specify more axes to np.tensordot in order to remove some of the loop nesting? In general looping and summing in pure python is slower than an equivalent call to a numpy function.