Apply interpolation about geowombat HOT 4 CLOSED

class interpolate_nan(gw.TimeModule):
    def __init__(self, missing_value=None, interp_type="linear", count=1):
        super(interpolate_nan, self).__init__()
        self.missing_value = missing_value
        self.interp_type = interp_type
        # Overrides the default output band count
        self.count = count

    def calculate(self, array):
        # check if missing_value is not None and not np.nan
        if self.missing_value is not None:
            if not np.isnan(self.missing_value):
                array = jnp.where(array == self.missing_value, np.NaN, array)
            if self.interp_type == "linear":
                array = np.apply_along_axis(_interpolate_nans_linear, axis=0, arr=array)
        # Return the interpolated array (3d -> time/bands x height x width)
        # If the array is (time x 1 x height x width) then squeeze to 3d
        return array.squeeze()

with gw.series(
    files,
    nodata=9999,
) as src:
    src.apply(
        func=interpolate_nan(
            missing_value=0, 
            # not sure if your output length matches your input file length
            # whatever your case is, this is where you define the output band count
            count=len(src.filenames)
        ),
        outfile=f"/home/mmann1123/Downloads/test.tif",
        num_workers=5,
        # Note that this is the band, or bands, to read
        bands=1,
    )

with gw.open(files, chunks={'time': -1}) as src:
    interp = (
        src.interpolate_na(dim='time', method='linear', fill_value='extrapolate')
        .bfill(dim='time')
        .ffill(dim='time')
        # Interpolate to new grid
        #.interp(time=smooth_range, method='slinear')
    )