How to use the verde.Trend function in verde

to remove the clear trend from our Texas temperature dataset
(:func:`verde.datasets.fetch_texas_wind`).
"""
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import verde as vd

# Load the Texas wind and temperature data as a pandas.DataFrame
data = vd.datasets.fetch_texas_wind()
print("Original data:")
print(data.head())

# Fit a 1st degree 2D polynomial to the data
coordinates = (data.longitude, data.latitude)
trend = vd.Trend(degree=1).fit(coordinates, data.air_temperature_c)
print("\nTrend estimator:", trend)

# Add the estimated trend and the residual data to the DataFrame
data["trend"] = trend.predict(coordinates)
data["residual"] = data.air_temperature_c - data.trend
print("\nUpdated DataFrame:")
print(data.head())


# Make a function to plot the data using the same colorbar
def plot_data(column, i, title):
    "Plot the column from the DataFrame in the ith subplot"
    crs = ccrs.PlateCarree()
    ax = plt.subplot(2, 2, i, projection=ccrs.Mercator())
    ax.set_title(title)
    # Set vmin and vmax to the extremes of the original data

# Trends
# ------
#
# Trends can't handle vector data automatically, so you can't pass
# ``data=(data.velocity_east, data.velocity_north)`` to :meth:`verde.Trend.fit`. To get
# around that, you can use the :class:`verde.Vector` class to create multi-component
# estimators and gridders from single component ones.
#
# :class:`~verde.Vector` takes an estimator/gridder for each data component and
# implements the :ref:`gridder interface ` for vector data, fitting
# each estimator/gridder given to a different component of the data.
#
# For example, to fit a trend to our GPS velocities, we need to make a 2-component
# vector trend:

trend = vd.Vector([vd.Trend(4), vd.Trend(1)])
print(trend)

########################################################################################
# We can use the ``trend`` as if it were a regular :class:`verde.Trend` but passing in
# 2-component data to fit. This will fit each data component to a different
# :class:`verde.Trend`.

trend.fit(
    coordinates=proj_coords,
    data=(data.velocity_east, data.velocity_north),
    weights=(1 / data.std_east ** 2, 1 / data.std_north ** 2),
)

########################################################################################
# Each estimator can be accessed through the ``components`` attribute:

plt.scatter(
    data.longitude,
    data.latitude,
    c=data.air_temperature_c,
    s=100,
    cmap="plasma",
    transform=ccrs.PlateCarree(),
)
plt.colorbar().set_label("Air temperature (C)")
vd.datasets.setup_texas_wind_map(ax)
plt.show()

########################################################################################
# We can estimate the polynomial coefficients for this trend:

trend = vd.Trend(degree=1).fit(coordinates, data.air_temperature_c)
print(trend.coef_)

########################################################################################
# More importantly, we can predict the trend values and remove them from our data:

trend_values = trend.predict(coordinates)
residuals = data.air_temperature_c - trend_values

fig, axes = plt.subplots(
    1, 2, figsize=(10, 6), subplot_kw=dict(projection=ccrs.Mercator())
)

ax = axes[0]
ax.set_title("Trend")
tmp = ax.scatter(
    data.longitude,

# Load the Rio de Janeiro total field magnetic anomaly data
data = vd.datasets.fetch_rio_magnetic()
region = vd.get_region((data.longitude, data.latitude))

# Create a projection for the data using pyproj so that we can use it as input for the
# gridder. We'll set the latitude of true scale to the mean latitude of the data.
projection = pyproj.Proj(proj="merc", lat_ts=data.latitude.mean())

# Create a chain that fits a 2nd degree trend, decimates the residuals using a blocked
# mean to avoid aliasing, and then fits a standard gridder to the residuals. The spacing
# for the blocked mean will be 0.5 arc-minutes (approximately converted to meters).
spacing = 0.5 / 60
chain = vd.Chain(
    [
        ("trend", vd.Trend(degree=2)),
        ("reduce", vd.BlockReduce(np.mean, spacing * 111e3)),
        ("spline", vd.Spline(damping=1e-8, mindist=50e3)),
    ]
)
print("Chained estimator:", chain)
# Calling 'fit' will automatically run the data through the chain
chain.fit(
    projection(data.longitude.values, data.latitude.values), data.total_field_anomaly_nt
)

# Each component of the chain can be accessed separately using the 'named_steps'
# attribute
grid_trend = chain.named_steps["trend"].grid()
print("\nTrend grid:")
print(grid_trend)

"""
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import numpy as np
import verde as vd


# Fetch the GPS data from the U.S. West coast. The data has a strong trend toward the
# North-West because of the relative movement along the San Andreas Fault System.
data = vd.datasets.fetch_california_gps()

# We'll fit a degree 2 trend on both the East and North components and weight the data
# using the inverse of the variance of each component.
# Note: Never use [Trend(...)]*2 as an argument to Vector. This creates references
# to the same Trend instance and will mess up the fitting.
trend = vd.Vector([vd.Trend(degree=2) for i in range(2)])
weights = vd.variance_to_weights((data.std_east ** 2, data.std_north ** 2))
trend.fit(
    coordinates=(data.longitude, data.latitude),
    data=(data.velocity_east, data.velocity_north),
    weights=weights,
)
print("Vector trend estimator:", trend)

# The separate Trend objects for each component can be accessed through the 'components'
# attribute. You could grid them individually if you wanted.
print("East component trend:", trend.components[0])
print("East trend coefficients:", trend.components[0].coef_)
print("North component trend:", trend.components[1])
print("North trend coefficients:", trend.components[1].coef_)

# We can make a grid with both trend components as data variables

random_state=1,
)

########################################################################################
# Now we can make a 2-component spline. Since :class:`verde.Vector` implements
# ``fit``, ``predict``, and ``filter``, we can use it in a :class:`verde.Chain` to build
# a pipeline.
#
# We need to use a bit of damping so that the weights can be taken into account. Splines
# without damping provide a perfect fit to the data and ignore the weights as a
# consequence.

chain = vd.Chain(
    [
        ("mean", vd.BlockMean(spacing=spacing * 111e3, uncertainty=True)),
        ("trend", vd.Vector([vd.Trend(1), vd.Trend(1)])),
        ("spline", vd.Vector([vd.Spline(damping=1e-10), vd.Spline(damping=1e-10)])),
    ]
)
print(chain)

########################################################################################
#
# .. warning::
#
#     Never generate the component gridders with ``[vd.Spline()]*2``. This will result
#     in each component being a represented by **the same Spline object**, causing
#     problems when trying to fit it to different components.
#
# Fitting the spline and gridding is exactly the same as what we've done before.

chain.fit(*train)

"""
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import numpy as np
import verde as vd


# Fetch the GPS data from the U.S. West coast. The data has a strong trend toward the
# North-West because of the relative movement along the San Andreas Fault System.
data = vd.datasets.fetch_california_gps()

# We'll fit a degree 2 trend on both the East and North components and weight the data
# using the inverse of the variance of each component.
# Note: Never use [Trend(...)]*2 as an argument to Vector. This creates references
# to the same Trend instance and will mess up the fitting.
trend = vd.Vector([vd.Trend(degree=2) for i in range(2)])
weights = vd.variance_to_weights((data.std_east ** 2, data.std_north ** 2))
trend.fit(
    coordinates=(data.longitude, data.latitude),
    data=(data.velocity_east, data.velocity_north),
    weights=weights,
)
print("Vector trend estimator:", trend)

# The separate Trend objects for each component can be accessed through the 'components'
# attribute. You could grid them individually if you wanted.
print("East component trend:", trend.components[0])
print("East trend coefficients:", trend.components[0].coef_)
print("North component trend:", trend.components[1])
print("North trend coefficients:", trend.components[1].coef_)

# We can make a grid with both trend components as data variables

# Split the data into a training and testing set. We'll fit the gridder on the training
# set and use the testing set to evaluate how well the gridder is performing.
train, test = vd.train_test_split(
    projection(*coordinates), (data.velocity_east, data.velocity_north), random_state=0
)

# We'll make a 15 arc-minute grid in the end.
spacing = 15 / 60

# Chain together a blocked mean to avoid aliasing, a polynomial trend to take care of
# the increase toward the coast, and finally the vector gridder using Poisson's ratio
# 0.5 to couple the two horizontal components.
chain = vd.Chain(
    [
        ("mean", vd.BlockReduce(np.mean, spacing * 111e3)),
        ("trend", vd.Vector([vd.Trend(degree=1) for i in range(2)])),
        ("spline", vd.VectorSpline2D(poisson=0.5, mindist=10e3)),
    ]
)
# Fit on the training data
chain.fit(*train)
# And score on the testing data. The best possible score is 1, meaning a perfect
# prediction of the test data.
score = chain.score(*test)
print("Cross-validation R^2 score: {:.2f}".format(score))

# Interpolate our horizontal GPS velocities onto a regular geographic grid and mask the
# data that are far from the observation points
grid_full = chain.grid(
    region, spacing=spacing, projection=projection, dims=["latitude", "longitude"]
)
grid = vd.distance_mask(

train, test = vd.train_test_split(
    projection(*coordinates),
    (data.wind_speed_east_knots, data.wind_speed_north_knots),
    random_state=2,
)

# We'll make a 20 arc-minute grid
spacing = 20 / 60

# Chain together a blocked mean to avoid aliasing, a polynomial trend (Spline usually
# requires de-trended data), and finally a Spline for each component. Notice that
# BlockReduce can work on multicomponent data without the use of Vector.
chain = vd.Chain(
    [
        ("mean", vd.BlockReduce(np.mean, spacing * 111e3)),
        ("trend", vd.Vector([vd.Trend(degree=1) for i in range(2)])),
        (
            "spline",
            vd.Vector([vd.Spline(damping=1e-10, mindist=500e3) for i in range(2)]),
        ),
    ]
)
print(chain)

# Fit on the training data
chain.fit(*train)
# And score on the testing data. The best possible score is 1, meaning a perfect
# prediction of the test data.
score = chain.score(*test)
print("Cross-validation R^2 score: {:.2f}".format(score))

# Interpolate the wind speed onto a regular geographic grid and mask the data that are