Bayesian inversion to produce estimates using a state-space time-series model

Hi,

I’m having difficulty getting pymc3 to work when inverting the outputs of an RTM. Please could someone offer me some advice?

Ideally I’d like to take a time series of reflectance values and get a time series of associated crop traits - by using the prosail model in combination with a state-space time series model and bayesian model.

import numpy as np
import pandas as pd
import pymc3 as pm
import theano.tensor as tt
import arviz as az
import prosail
from theano.compile.ops import as_op

# Preprocessing: normalize the reflectance data reflectance has shape (,12)
reflectance_data_normalized = (reflectance - reflectance.mean()) / reflectance.std()

@as_op(itypes=[tt.dscalar, tt.dscalar, tt.dscalar, tt.dscalar, tt.dscalar, tt.dscalar], otypes=[tt.dvector])
def prosail_model(n, cab, car, cbrown, cw, cm):
    # ProSail model with parameters: n, cab, car, cbrown, cw, cm
    # Returns simulated reflectances
    reflectances = []
    reflectance = prosail.run_prosail(n, cab, car, cbrown, cw, cm, w,
                                          lidfa=-1.0, hspot=0.01, tts=30., tto=10., psi=0., factor='SDR', rsoil=0.5, psoil=0.5, prospect_version='D')
    return reflectances

# Model block
with pm.Model() as model:
    # Priors for ProSail parameters
    n_prior = pm.Uniform("n", lower=0.0, upper=1.0)
    cab_prior = pm.Uniform("cab", lower=0.0, upper=100.0)
    car_prior = pm.Uniform("car", lower=0.0, upper=100.0)
    cbrown_prior = pm.Uniform("cbrown", lower=0.0, upper=100.0)
    cw_prior = pm.Uniform("cw", lower=0.0, upper=0.5)
    cm_prior = pm.Uniform("cm", lower=0.0, upper=0.5)

    # Combine ProSail parameters
    prosail_prior = tt.stack([n_prior, cab_prior, car_prior, cbrown_prior, cw_prior, cm_prior], axis=-1)

    # Calculate ProSail model output for each time step
    prosail_output = pm.Deterministic("prosail_output", calculate_prosail(prosail_prior))

    # Likelihood - assuming normally distributed noise
    sigma_obs = pm.HalfNormal("sigma_obs", sigma=0.1)
    obs = pm.Normal("obs", mu=prosail_output, sigma=sigma_obs, observed=reflectance_data_normalized)

    # Sample
    trace = pm.sample(2000, tune=1000, target_accept=0.9)

# Plot traces
az.plot_trace(trace)

My error is “TypeError: Expected an array-like object, but found a Variable: maybe you are trying to call a function on a (possibly shared) variable instead of a numeric array?” because I am struggling to get the prosail scipy functions to cooperate with pymc3.

You can’t just use generic numpy/scipy code with PyMC. As you discovered, these functions expect numeric values to compute on, but a PyMC graph is comprised of symbolic pytensor Variables. Thus the error. In some cases it works, but it’s a happy accident. The computational backend for PyMC is pytensor – see here for some discussion on how it all works.

You were on the right track when using as_op – you will need to wrap your prosail model in a Pytensor Op. Well first you need to upgrade your setup if possible – pymc3 and theano are several years out of date at this point. After that, you can check out this tutorial about wrapping arbitrary numpy code in a custom Ops for guidance on how to proceed.

Thanks - I’ve had another go at this and come up with this code:

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pymc as pm
import pytensor
import pytensor.tensor as pt
import prosail  # Importing the ProSail package

from pytensor.graph import Apply, Op
from scipy.optimize import approx_fprime

def prosail_model(n, cab, car, cbrown, cw, cm, lai, tto, tts):
    # ProSail model with parameters: n, cab, car, cbrown, cw, cm, lai, tto, tts
    reflectance = prosail.run_prosail(n=n,
                                      cab=cab,
                                      car=car,
                                      cbrown=cbrown,
                                      cw=cw,
                                      cm=cm,
                                      lai=1.5,
                                      lidfa=-1.0,
                                      hspot=0.01,
                                      tts=10.,
                                      tto=30.,
                                      psi=0.,
                                      prospect_version="D",
                                      factor='SDR', 
                                      rsoil=.5, psoil=.5)
    # reflectances2 = Prosail2S2("s2_response.csv", reflectance)
    return reflectance[:13]

def prosail_loglike(n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, data):
    for param in (n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, data):
        if not isinstance(param, (float, np.ndarray)):
            raise TypeError(f"Invalid input type to loglike: {type(param)}")

    model = prosail_model(n, cab, car, cbrown, cw, cm, lai, tto, tts)
    return -0.5 * ((data - model) / sigma) ** 2 - np.log(np.sqrt(2 * np.pi)) - np.log(sigma)

# Define a PyTensor Op for our ProSail likelihood function
class LogLike(Op):
    def make_node(self, n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, data) -> Apply:
        n = pt.as_tensor(n)
        cab = pt.as_tensor(cab)
        car = pt.as_tensor(car)
        cbrown = pt.as_tensor(cbrown)
        cw = pt.as_tensor(cw)
        cm = pt.as_tensor(cm)
        lai = pt.as_tensor(lai)
        tto = pt.as_tensor(tto)
        tts = pt.as_tensor(tts)
        sigma = pt.as_tensor(sigma)
        data = pt.as_tensor(data)
        
        inputs = [n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, data]
        outputs = [data.type()]

        return Apply(self, inputs, outputs)

    def perform(self, node: Apply, inputs: list[np.ndarray], outputs: list[list[None]]) -> None:
        n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, data = inputs

        loglike_eval = prosail_loglike(n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, data)

        outputs[0][0] = np.asarray(loglike_eval)
        
# Set up our data
N = 13  # number of data points
sigma = 1.0  # standard deviation of noise

# True parameter values for ProSail
ntrue = 1.5
cabtrue = 40
cartrue = 10
cbrowntrue = 1
cwtrue = 0.015
cmtrue = 0.009
laitrue = 1.5
ttotrue = 10
ttstrue = 30

# Make data
truemodel = prosail_model(ntrue, cabtrue, cartrue, cbrowntrue, cwtrue, cmtrue, laitrue, ttotrue, ttstrue)
rng = np.random.default_rng(716743)
# data = reflectance[5,:]/10000 
data = sigma * rng.normal(size=N) + truemodel

# Create our Op
loglike_op = LogLike()

# test_out = loglike_op(ntrue, cabtrue, cartrue, cbrowntrue, cwtrue, cmtrue, sigma, data)

def custom_dist_loglike(data, n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma):
    return loglike_op(n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, data)

# Time series setup
T = 10  # Number of time points
times = np.arange(T)

# Generate synthetic time series data for true parameters
n_true = np.cumsum(rng.normal(0, 0.1, T)) + ntrue
cab_true = np.cumsum(rng.normal(0, 5, T)) + cabtrue
car_true = np.cumsum(rng.normal(0, 2, T)) + cartrue
cbrown_true = np.cumsum(rng.normal(0, 0.1, T)) + cbrowntrue
cw_true = np.cumsum(rng.normal(0, 0.001, T)) + cwtrue
cm_true = np.cumsum(rng.normal(0, 0.001, T)) + cmtrue
lai_true = np.cumsum(rng.normal(0, 0.1, T)) + laitrue
tto_true = np.cumsum(rng.normal(0, 0.1, T)) + ttotrue
tts_true = np.cumsum(rng.normal(0, 0.1, T)) + ttstrue

# Simulate reflectance data for each time point
data = np.array([prosail_model(n, cab, car, cbrown, cw, cm, lai, tto, tts) for n, cab, car, cbrown, cw, cm, lai, tto, tts in zip(n_true, cab_true, car_true, cbrown_true, cw_true, cm_true, lai_true, tto_true, tts_true)])
data += sigma * rng.normal(size=data.shape)

# Use PyMC to sample from the state-space model
with pm.Model() as state_space_model:
    # Gaussian random walk priors for ProSail parameters
    n = pm.GaussianRandomWalk("n", sigma=0.1, shape=T, init_dist=pm.Uniform.dist(1,5))
    cab = pm.GaussianRandomWalk("cab", sigma=5, shape=T, init_dist=pm.Uniform.dist(5,60))
    car = pm.GaussianRandomWalk("car", sigma=2, shape=T, init_dist=pm.Uniform.dist(5,20))
    cbrown = pm.GaussianRandomWalk("cbrown", sigma=0.1, shape=T, init_dist=pm.Uniform.dist(0,2))
    cw = pm.GaussianRandomWalk("cw", sigma=0.001, shape=T, init_dist=pm.Uniform.dist(0,0.02))
    cm = pm.GaussianRandomWalk("cm", sigma=0.001, shape=T, init_dist=pm.Uniform.dist(0,0.02))

    lai = pm.GaussianRandomWalk("lai", sigma=0.1, shape=T, init_dist=pm.Uniform.dist(0.5,5))
    tto = pm.GaussianRandomWalk("tto", sigma=0.1, shape=T, init_dist=pm.Uniform.dist(0,20))
    tts = pm.GaussianRandomWalk("tts", sigma=0.1, shape=T, init_dist=pm.Uniform.dist(20,60))

    resid_var = pm.Uniform("resid_var", lower=0.0, upper=0.1, initval=0.01)

    # Likelihood
    for t in range(T):
        pm.CustomDist(
            f"likelihood_{t}", 
            n[t], cab[t], car[t], cbrown[t], cw[t], cm[t], lai[t], tto[t], tts[t], resid_var, 
            observed=data[t], logp=custom_dist_loglike
        )

    # Sampling
    idata_state_space = pm.sample(100, tune=10, step=pm.DEMetropolisZ())

# Plot the traces, space out the plots for better visibility
az.plot_trace(idata_state_space)
plt.subplots_adjust(hspace=0.5)
plt.show()

Unfortunately I have a few problems remaining with this… Any help would be appreciated.

  • In the prosail call I set lai = 1.5. This is because introducing it as a parameter stops the code working. It suggests the argument types are incorrect but that doesn’t make any sense considering what I’m trying to do. Is this a bug?
  • I am using a random walk time series approach to combining inversions through time in order to generate more constrained results compared with the inversion posterior distributions by themselves. Is this the right way to do it? Your tutorials on state space models have no bearing on what I’m trying to do - suggesting I’m missing something here…

Many thanks

I don’t know anything about prosail, so I can’t comment on that. I’m also not sure what you mean by an “inversion posterior”.

Typically in time-varying statespace setups, you would create states for the time-varying parameters and insert the data into the design matrix. See here for example. I don’t think there’s anything wrong with your approach – I guess it will lead to a time varying transition matrix? – but if you’re having identification problems it might be useful to try the other way.

Thanks @jessegrabowski. I’m having problems inverting prosail as just a single point bayesian inversion problem, so I’ll work on that first before moving onto the time series approach.
As long as my example code looks like a vaguely sensible approach to “model something in time so it has more sensible results that just modelling each point in time individually”. If it helps, this is the R code that I am trying to port to python: RSE-00511/Code/RTM_MCMC_V2.Rmd at main · DongchenZ/RSE-00511 · GitHub

Regarding my first bullet point, I am getting a confusing TypeError when I include lai. If you are able to run my code snippet yourself and pass comment on what the cause could be that would be greatly appreciated… This is my MWE:

import numpy as np
import pymc as pm
import pytensor.tensor as pt
from pytensor.graph import Op, Apply
import arviz as az
import prosail

# Define the new model function with additional parameters
def my_model(n, cab, car, cbrown, cw, cm, lai, tto, tts, x):
    reflectance = prosail.run_prosail(n=n, cab=cab, car=car, cbrown=cbrown, cw=cw, cm=cm, lai=1.5, tto=10., tts=30., lidfa=-1.0, hspot=0.01, psi=0., prospect_version="D", factor='SDR', rsoil=.5, psoil=.5)
    return reflectance #Prosail2S2("s2_response.csv", reflectance)
    # Example model function - you can replace this with your specific model
    return (n * x + cab * x**2 + car * np.sin(x) + cbrown * np.cos(x) +
            cw * x**0.5 + cm * np.log1p(x) + lai * np.exp(-x) +
            tto * np.tan(x) + tts * np.tanh(x))

def my_loglike(n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x, data):
    for param in (n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x, data):
        if not isinstance(param, (float, np.ndarray)):
            raise TypeError(f"Invalid input type to loglike: {type(param)}")
    model = my_model(n, cab, car, cbrown, cw, cm, lai, tto, tts, x)
    return -0.5 * ((data - model) / sigma) ** 2 - np.log(np.sqrt(2 * np.pi)) - np.log(sigma)

class LogLike(Op):
    def make_node(self, n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x, data) -> Apply:
        n = pt.as_tensor(n)
        cab = pt.as_tensor(cab)
        car = pt.as_tensor(car)
        cbrown = pt.as_tensor(cbrown)
        cw = pt.as_tensor(cw)
        cm = pt.as_tensor(cm)
        lai = pt.as_tensor(lai, dtype=float)
        tto = pt.as_tensor(tto)
        tts = pt.as_tensor(tts)
        sigma = pt.as_tensor(sigma)
        x = pt.as_tensor(x)
        data = pt.as_tensor(data)

        inputs = [n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x, data]
        outputs = [data.type()]

        return Apply(self, inputs, outputs)

    def perform(self, node: Apply, inputs: list[np.ndarray], outputs: list[list[None]]) -> None:
        n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x, data = inputs
        loglike_eval = my_loglike(n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x, data)
        outputs[0][0] = np.asarray(loglike_eval)

def custom_dist_loglike(data, n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x):
    return loglike_op(n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x, data)

N = 2101
sigma = 1.0
x = np.linspace(0.0, 9.0, N)

# True parameter values for testing
ntrue = 0.4
cabtrue = 0.1
cartrue = 0.2
cbrowntrue = 0.3
cwtrue = 0.05
cmtrue = 0.07
laitrue = 3.0
ttotrue = 0.15
ttstrue = 0.25

truemodel = my_model(ntrue, cabtrue, cartrue, cbrowntrue, cwtrue, cmtrue, laitrue, ttotrue, ttstrue, x)

rng = np.random.default_rng(716743)
data = sigma * rng.normal(size=N) + truemodel

loglike_op = LogLike()

test_out = loglike_op(ntrue, cabtrue, cartrue, cbrowntrue, cwtrue, cmtrue, laitrue, ttotrue, ttstrue, sigma, x, data)

with pm.Model() as no_grad_model:
    # n = pm.Uniform("n", lower=-10, upper=10.0, initval=ntrue)
    # cab = pm.Uniform("cab", lower=-10, upper=60.0, initval=cabtrue)
    # car = pm.Uniform("car", lower=-10.0, upper=20.0, initval=cartrue)
    # cbrown = pm.Uniform("cbrown", lower=-10.0, upper=10.0, initval=cbrowntrue)
    # cw = pm.Uniform("cw", lower=-10.0, upper=10.0, initval=cwtrue)
    # cm = pm.Uniform("cm", lower=-10.0, upper=10.0, initval=cmtrue)
    # lai = pm.Uniform("lai", lower=-10.0, upper=10.0, initval=laitrue)
    # tto = pm.Uniform("tto", lower=-10.0, upper=20.0, initval=ttotrue)
    # tts = pm.Uniform("tts", lower=-10.0, upper=60.0, initval=ttstrue)

    n = pm.Uniform("n", lower=-1.0, upper=5.0, initval=ntrue)
    cab = pm.Uniform("cab", lower=-5, upper=60, initval=cabtrue)
    car = pm.Uniform("car", lower=-5, upper=20, initval=cartrue)
    cbrown = pm.Uniform("cbrown", lower=0, upper=2, initval=cbrowntrue)
    cw = pm.Uniform("cw", lower=-10.0, upper=10, initval=cwtrue)
    cm = pm.Uniform("cm", lower=0, upper=0.2, initval=cmtrue)
    lai = pm.Uniform("lai", lower=0.5, upper=5.0, initval=laitrue)
    tto = pm.Uniform("tto", lower=0.0, upper=20.0, initval=ttotrue)
    tts = pm.Uniform("tts", lower=0, upper=60.0, initval=ttstrue)

    # sigma = pm.Uniform("sigma", lower=0, upper=3, initval=sigma)

    likelihood = pm.CustomDist(
        "likelihood", n, cab, car, cbrown, cw, cm, lai, tto, tts, sigma, x, observed=data, logp=custom_dist_loglike
    )

ip = no_grad_model.initial_point()
ip

no_grad_model.compile_logp(vars=[likelihood], sum=False)(ip)

try:
    no_grad_model.compile_dlogp()
except Exception as exc:
    print(type(exc))

import warnings
with warnings.catch_warnings():
    warnings.simplefilter("ignore")

    with no_grad_model:
        idata_no_grad = pm.sample(10000, step=pm.DEMetropolisZ())

az.plot_trace(idata_no_grad, lines=[
    ("n", {}, ntrue), ("cab", {}, cabtrue), 
    ("car", {}, cartrue), ("cbrown", {}, cbrowntrue), 
    ("cw", {}, cwtrue), ("cm", {}, cmtrue), 
    ("lai", {}, laitrue), ("tto", {}, ttotrue), 
    ("tts", {}, ttstrue)
])

az.summary(idata_no_grad, round_to=2)