Error in so3_rotation (Jd matrix size) with custom data

[SkirOwen]

I'm trying to make the equivariance_plot script work on my own weather data of size (721, 1440).
When it runs the so3_rotation function I have the following error:

Traceback (most recent call last):
  File "/home/owen/anaconda3/envs/cupy-test/lib/python3.8/site-packages/IPython/core/interactiveshell.py", line 3441, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-2-678a6fff5083>", line 1, in <module>
    runfile('/home/owen/PycharmProjects/cranfield/irp/ai.py', wdir='/home/owen/PycharmProjects/cranfield/irp')
  File "/home/owen/.local/share/JetBrains/Toolbox/apps/PyCharm-P/ch-0/211.7628.24/plugins/python/helpers/pydev/_pydev_bundle/pydev_umd.py", line 197, in runfile
    pydev_imports.execfile(filename, global_vars, local_vars)  # execute the script
  File "/home/owen/.local/share/JetBrains/Toolbox/apps/PyCharm-P/ch-0/211.7628.24/plugins/python/helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "/home/owen/PycharmProjects/cranfield/irp/ai.py", line 109, in <module>
    main(date, cycle)
  File "/home/owen/PycharmProjects/cranfield/irp/ai.py", line 80, in main
    y1 = phi(s2_rotation(x, *abc))
  File "/home/owen/PycharmProjects/cranfield/irp/ai.py", line 15, in s2_rotation
    x = so3_rotation(x.view(*x.size(), 1).expand(*x.size(), x.size(-1)), a, b, c)
  File "/home/owen/PycharmProjects/cranfield/irp/s2cnn/soft/so3_rotation.py", line 18, in so3_rotation
    Us = _setup_so3_rotation(b, alpha, beta, gamma, device_type=x.device.type, device_index=x.device.index)
  File "/home/owen/PycharmProjects/cranfield/irp/s2cnn/soft/so3_rotation.py", line 67, in _setup_so3_rotation
    Us = __setup_so3_rotation(b, alpha, beta, gamma)
  File "/home/owen/PycharmProjects/cranfield/irp/s2cnn/utils/decorator.py", line 97, in wrapper
    result = func(*args)
  File "/home/owen/PycharmProjects/cranfield/irp/s2cnn/soft/so3_rotation.py", line 55, in __setup_so3_rotation
    Us = [wigner_D_matrix(l, alpha, beta, gamma,
  File "/home/owen/PycharmProjects/cranfield/irp/s2cnn/soft/so3_rotation.py", line 55, in <listcomp>
    Us = [wigner_D_matrix(l, alpha, beta, gamma,
  File "/home/owen/anaconda3/envs/cupy-test/lib/python3.8/site-packages/lie_learn/representations/SO3/wigner_d.py", line 63, in wigner_D_matrix
    D = rot_mat(alpha=alpha, beta=beta, gamma=gamma, l=l, J=Jd[l])
IndexError: index 151 is out of bounds for axis 0 with size 151

After trying to debug it, I have found that the issue is with the Jd matrix of size 151, defined in lie_learn when calling the line 55 in s2cnn/soft/so3_rotation.py

Us = [wigner_D_matrix(l, alpha, beta, gamma,
                          field='complex', normalization='quantum', order='centered', condon_shortley='cs')
          for l in range(b)]

b is set as b = x.size()[-1] // 2 so 720 in my case

the winger_D_matrix calls the rot_mat function which uses Jd[l], but since Jd is size 151, it throws the error.

Is there a way to redefine the Jd matrix or do I need to use a different function?