Term.get_coefficients fails for jax implementation when built from source

[bmorris3]

Hi @dfm,

Today I've built celerite2 from source following the recommendations on the install docs. I'm trying to do something simple, like this

from celerite2.jax import terms

sho = terms.SHOTerm(S0=1, w0=3000, Q=0.6)
sho.get_coefficients()

but I'm getting the following error

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
Cell In [1], line 4
      1 from celerite2.jax import terms
      3 sho = terms.SHOTerm(S0=1, w0=3000, Q=0.6)
----> 4 sho.get_coefficients()

File ~/git/celerite2/python/celerite2/jax/terms.py:36, in Term.get_coefficients(self)
     35 def get_coefficients(self):
---> 36     raise NotImplementedError("subclasses must implement this method")

NotImplementedError: subclasses must implement this method

At first I thought this could be an accident of the multiple SHOTerm implementations, for example, here

celerite2/python/celerite2/jax/terms.py

Lines 473 to 478 in e75dd45

	def SHOTerm(*args, **kwargs):
	over = OverdampedSHOTerm(*args, **kwargs)
	under = UnderdampedSHOTerm(*args, **kwargs)
	if over.Q < 0.5:
	return over
	return under

and here

celerite2/python/celerite2/jax/terms.py

Line 481 in e75dd45

class SHOTerm(Term):

but commenting the first one out doesn't solve the problem.

Any ideas? Thanks!

[bmorris3]

Writing an issue is always the most clarifying exercise.

I was hitting this error because I was using TermConvolution, which uses this line:

celerite2/python/celerite2/jax/terms.py

Lines 301 to 310 in e75dd45

	class TermConvolution(Term):
	def __init__(self, term, delta):
	self.delta = np.float64(delta)
	try:
	self.coefficients = term.get_coefficients()
	except NotImplementedError:
	raise TypeError(
	"Term operations can only be performed on terms that provide "
	"coefficients"
	)

but the jax implementation of SHOTerm has under/overdamped versions of the coefficients defined separately. Should TermConvolution be modified accordingly?

[dfm]

Great question! This is an issue introduced by this PR: #68

I should probably spend some time thinking about how to fix this properly, but one option for the short term (it'll take a minor? performance hit) would be something like:

import jax
import jax.numpy as jnp
from celerite2.jax import terms

def custom_sho_get_coeffs(self):
    ar, cr = self.get_overdamped_coefficients()
    ac, bc, cc, dc = self.get_underdamped_coefficients()
    cond = jnp.less(self.Q, 0.5)
    selectr = lambda x: jax.lax.cond(cond, lambda y: y, jnp.zeros_like, operand=x)
    selectc = lambda x: jax.lax.cond(cond, jnp.zeros_like, lambda y: y, operand=x)
    return selectr(ar), selectr(cr), selectc(ac), selectc(bc), selectc(cc), selectc(dc)

terms.SHOTerm.get_coefficients = custom_sho_get_coeffs
sho = terms.SHOTerm(S0=1, w0=3000, Q=0.6)
sho.get_coefficients()

[bmorris3]

I can confirm the fix above works for me for now. Thanks Dan!