{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "545bcee1-7870-425c-9a86-8843ebfd8d87",
   "metadata": {},
   "source": [
    "# Compute SST fitting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7bb1fa4b-d8b6-4287-bc79-feb9f2bf6fd7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as N\n",
    "import xarray as xr\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import sys\n",
    "import dask\n",
    "from os import path\n",
    "from dask.distributed import Client, LocalCluster, progress\n",
    "import time\n",
    "from tools import my_percentile\n",
    "from scipy.optimize import least_squares,curve_fit\n",
    "import dask_hpcconfig\n",
    "from dask_jobqueue import PBSCluster\n",
    "import glob"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7e69237-8907-4f3a-896e-560255b954c6",
   "metadata": {},
   "source": [
    "### Set Dask workers in local mode (use all 28 cpus on the nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "293e8282-f4c5-42e5-a5ec-602afa308123",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/user/mcaillau/proxy/8787/status\n"
     ]
    }
   ],
   "source": [
    "#LOCAL\n",
    "overrides = {\"cluster.n_workers\": 28,\"cluster.threads_per_worker\":1}\n",
    "cluster = dask_hpcconfig.cluster(\"datarmor-local\",**overrides)\n",
    "print(cluster.dashboard_link)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7d54b8c0-7f90-4d6f-b55a-253be3db0568",
   "metadata": {},
   "outputs": [],
   "source": [
    "# explicitly connect to the cluster we just created\n",
    "client = Client(cluster)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfcf3a36-70b0-4134-b566-02bced132407",
   "metadata": {},
   "source": [
    "### Read croco grid "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "56a0cd1d-7ad6-4d76-88e8-117fa5429987",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#read grid\n",
    "ds_grid=xr.open_dataset('/home/shom_simuref/CROCO/ODC/CONFIGS/MEDITERRANEE_GLOBALE/CROCO_FILES/test2.nc')\n",
    "coord_dict={\"xi_rho\":\"X\",\"eta_rho\":\"Y\",\"xi_u\":\"X_U\",\"eta_v\":\"Y_V\"}\n",
    "ds_grid=ds_grid.assign_coords({\"X\":ds_grid.lon_rho[0,:], \"Y\":ds_grid.lat_rho[:,0]})\n",
    "ds_grid=ds_grid.swap_dims(coord_dict)\n",
    "mask=ds_grid.mask_rho"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7eed301a-3cc5-439f-b47c-82ac00ab0a94",
   "metadata": {},
   "source": [
    "### Read stats netcdf file from another script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ea0a7d74-13a3-422e-a219-1cc676143524",
   "metadata": {},
   "outputs": [],
   "source": [
    "#read croco\n",
    "stat_dir='/home/shom_simuref/CROCO/ODC/POSTPROC/SST/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ae5157d2-511e-465e-a2a2-d29b72eeade7",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds=xr.open_dataset(f'{stat_dir}/mean_sst.nc')\n",
    "sst_mean=ds.temp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5c09e47d-3ef7-44a7-9152-22527c9f2d04",
   "metadata": {},
   "outputs": [],
   "source": [
    "#remove last index to have only 12 months\n",
    "sst_mean=sst_mean.isel(time=slice(0,-1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99135484-103b-4973-bbb0-86eac56927ca",
   "metadata": {},
   "source": [
    "### Define fitting function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "11d39b72-4126-4a8c-b66a-c2179e0e292c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "n_months = sst_mean.shape[0]      # number of months/points of the time-series\n",
    "\n",
    "### curve fit model\n",
    "def model_curve(t,a,b,c) :\n",
    "    sst_cycle_model = a*N.cos(2*N.pi*1/12*t-b)+c\n",
    "    #sst_cycle_model = a*N.cos(2*N.pi*1/12*t-b)+c+d*t\n",
    "    return sst_cycle_model\n",
    "\n",
    "### least square model (for comparison only) ####\n",
    "def model_1freq_trend(x) :\n",
    "    t = N.arange(0,n_months)\n",
    "    sst_cycle_model = x[0]*N.cos(2*N.pi*1/12*t-x[1])+x[2]\n",
    "    return sst_cycle_model\n",
    "\n",
    "#cost function \n",
    "def cost_function_1freq(x,sst_signal):\n",
    "    sst_cycle_model = model_1freq_trend(x)\n",
    "    return (sst_cycle_model-sst_signal)\n",
    "\n",
    "# fitting with least square function\n",
    "def optimize_1freq(x_0,sst_signal,loss):\n",
    "    res_ar   = least_squares(cost_function_1freq, x_0,args=([sst_signal]),method='trf',loss=loss, f_scale=1.0)\n",
    "    return res_ar.x,res_ar.cost\n",
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "354bbdfd-5edd-46c8-93d3-9608d026d05e",
   "metadata": {},
   "source": [
    "###  CROCO FITTING"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "355ca8ae-0492-48c5-917e-63528c84d1dd",
   "metadata": {},
   "source": [
    "### Pre process of data before fitting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d0893fcf-0fe6-47bf-af70-29318f000786",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "calc init parameters\n",
      "(12.0451565, 0, 20.337656)\n"
     ]
    }
   ],
   "source": [
    "#least_square initialisation guess from one point \n",
    "print('calc init parameters')\n",
    "\n",
    "#set a new time vector\n",
    "xdata= N.arange(0,n_months)\n",
    " \n",
    "#assign new coordinates\n",
    "sst2=sst_mean.assign_coords(month=(\"time\", xdata))\n",
    "sst2=sst2.swap_dims({\"time\": \"month\"})\n",
    "\n",
    "#here we use small chunks because the calcul of curvefit in xarray seem not parallelized\n",
    "#it's more efficient to have a lot of chunks of the domain\n",
    "#month should not be chunked because the fitting is applied along this dimension\n",
    "sst2=sst2.chunk({\"month\":-1,\"Y\":100,\"X\":100})\n",
    "\n",
    "#set land values to Nan to avoid to fit theses values\n",
    "sst2=sst2.where(sst2>0)\n",
    "\n",
    "\n",
    "#set bounds for curvefit to avoid negtive temperature \n",
    "bounds={\"a\":(0,N.inf),\"b\":(-N.inf,N.inf),\"c\":(0,N.inf)}\n",
    "\n",
    "#set initial guess (not mandatory)\n",
    "signal=sst2[:,400,900].compute()\n",
    "x_0 = N.max(signal.data)-N.min(signal.data),0,N.mean(signal.data)\n",
    "p0={\"a\":x_0[0],\"b\":x_0[1],\"c\":x_0[2]}\n",
    "print(x_0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57568db8-0582-44d5-954b-d3a1b84b4c40",
   "metadata": {},
   "source": [
    "### define the curvefit model with xarray method \"curvefit\" (that calls scipy.curvefit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "35164149-187f-4d2b-969b-1f12130fc75e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 12 ms, sys: 0 ns, total: 12 ms\n",
      "Wall time: 14.1 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "sst_fit= sst2.curvefit(\n",
    "       coords=\"month\", func=model_curve,skipna=True,bounds=bounds)#,p0=p0)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbcf63a9-93fd-48a3-bb71-06e0096bdc0c",
   "metadata": {},
   "source": [
    "### Do the computation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "888d0cfa-a629-4e6d-ba7c-642380229dfa",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 54.2 s, sys: 8.98 s, total: 1min 3s\n",
      "Wall time: 4min 30s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "## 4min30 sur le client dask local avec 28 workers\n",
    "sst_fit=sst_fit.compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac3e8ebc-58cf-4473-acf2-420807b04c0a",
   "metadata": {},
   "source": [
    "### record data to netcdf file\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "286e15fe-c1f1-4b77-8cd5-009e6722f1cf",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'sst_fit' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Input \u001b[0;32mIn [1]\u001b[0m, in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m#record data to netcdf file\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m \u001b[43msst_fit\u001b[49m\u001b[38;5;241m.\u001b[39mto_netcdf(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mstat_dir\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m/sst_fit_croco.nc\u001b[39m\u001b[38;5;124m'\u001b[39m,mode\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mw\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'sst_fit' is not defined"
     ]
    }
   ],
   "source": [
    "sst_fit.to_netcdf(f'{stat_dir}/sst_fit_croco.nc',mode=\"w\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "988503f0-9c6e-43e2-a1bf-e86b67176656",
   "metadata": {},
   "source": [
    "### SEVIRI FITTING"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "baa79a64-8cbd-40d9-8f31-bc968be27bb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_sevi=xr.open_dataset(f'{stat_dir}/sst_sevi_mean.nc')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "6e5fadea-40c9-40d9-a42d-54a2ba132f3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "sst_sevi=ds_sevi.sea_surface_temperature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "61804b1a-b46e-4976-99e9-4584c07842b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "sst_sevi2=sst_sevi.assign_coords(month=(\"time\", xdata))\n",
    "sst_sevi2=sst_sevi2.swap_dims({\"time\": \"month\"})\n",
    "sst_sevi2-=273.15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "7522cf16-5107-4e0f-83d4-ebd081c8e90f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sst_sevi2=sst_sevi2.chunk({\"month\":-1,\"lat\":60,\"lon\":100})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "636640ca-5a87-413c-a353-853ab25584fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "not_null=sst_sevi2.notnull()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "4c3bfd97-6689-4554-8e93-4f2dc7807c12",
   "metadata": {},
   "outputs": [],
   "source": [
    "test=sst_sevi2.where(not_null.all(dim=\"month\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "89d5c694-7c00-4d6c-8cdd-6389048aba26",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.57 s, sys: 224 ms, total: 1.8 s\n",
      "Wall time: 5.3 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "sst_fit_sevi= test.curvefit(\n",
    "       coords=\"month\", func=model_curve,p0=p0).compute()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d5461f8-715a-462f-97fe-4d782e218584",
   "metadata": {},
   "source": [
    "### write results into a Netcdf file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "94ff192a-4c65-4579-a762-2f03a773d216",
   "metadata": {},
   "outputs": [],
   "source": [
    "sst_fit_sevi.to_netcdf(f'{stat_dir}/sst_fit_sevi.nc',mode=\"w\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d3d1393-4b23-47f5-a9d2-c73193649f71",
   "metadata": {},
   "source": [
    "### DEBUG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d53d8d48-1e46-4dcc-a240-b72437ce9c37",
   "metadata": {},
   "outputs": [],
   "source": [
    "#check curve fitting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a489a5c1-4ddb-499b-ae09-1df30e0cd993",
   "metadata": {},
   "outputs": [],
   "source": [
    "amp=sst_fit_sevi.curvefit_coefficients.sel(param=\"a\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7665cd0-3b86-4c51-9b91-db7bc3c45344",
   "metadata": {},
   "outputs": [],
   "source": [
    "N.argwhere((amp.data<0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb1c5637-60e3-4726-935c-60879c117d16",
   "metadata": {},
   "outputs": [],
   "source": [
    "signal = sst_sevi2[:,16,740]\n",
    "x_0 = N.max(signal.data)-N.min(signal.data),0,N.mean(signal.data)\n",
    "ydata=signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb0cf0c3-c471-4dbe-b673-dbb2bbec053d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#curve fit from scipy\n",
    "popt, pcov = curve_fit(model_curve, xdata, ydata, bounds=(0, [30, 365, 30]))#,p0=x_0) \n",
    "#curve from xarray\n",
    "coeffs=sst_fit_sevi.curvefit_coefficients[16,740].data\n",
    "#least square fit\n",
    "x,cost = optimize_1freq(x_0,ydata.data,loss=\"linear\")\n",
    "sig_fitted = model_1freq_trend(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec369062-5254-4834-b789-384f3e68a4da",
   "metadata": {},
   "outputs": [],
   "source": [
    "coeffs,popt,x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9ebf0a2-1364-437c-87b3-ec035723cd12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot\n",
    "fig,ax=plt.subplots(1,1,figsize=(10,8))\n",
    "ax.plot(xdata, ydata, 'b-', label='data')\n",
    "ax.plot(xdata,model_curve(xdata,*popt),'r-',label='scipy curve fit')\n",
    "ax.plot(xdata,model_curve(xdata,*coeffs),'g+',label='xarray wrapper curve fit')\n",
    "ax.plot(xdata,sig_fitted,'k*',markersize=2,label='least squares')\n",
    "plt.legend()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38347792-8391-491f-b094-804d388ce110",
   "metadata": {},
   "source": [
    "### PLOT THE RESULTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdc80f57-a087-430e-9e1a-d1e354bc704f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cartopy.crs as ccrs\n",
    "from cartopy.feature import ShapelyFeature\n",
    "import cartopy.feature as cfeature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f1eb6cb-ab25-4168-a0f7-136a768b4a37",
   "metadata": {},
   "outputs": [],
   "source": [
    "proj=ccrs.LambertConformal(central_latitude=38,central_longitude=15)\n",
    "lon_croco=sst_fit.X\n",
    "lat_croco=sst_fit.Y\n",
    "lon_sevi=sst_sevi2.lon\n",
    "lat_sevi=sst_sevi2.lat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c6fe9b5-7657-4ab5-a65c-d1ac7073d113",
   "metadata": {},
   "source": [
    "#### convert phase to days"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7e0f715-3640-45d8-8411-e209b528c0ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_phase(phase):\n",
    "    phase=phase/2*N.pi*365\n",
    "    arr=phase.data\n",
    "    arr[arr>365]-=365\n",
    "    phase[:]=arr[:]\n",
    "    print(phase.min(),phase.max())\n",
    "    return phase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10abc647-8ebb-4d74-9871-d6e9892cce72",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sst_croco=sst_fit.curvefit_coefficients\n",
    "phase_croco=sst_croco.sel(param=\"b\")\n",
    "phase_croco2=convert_phase(phase_croco)\n",
    "print(phase_croco2.min(),phase_croco2.max())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28f490f2-7a12-4193-ad89-f8179701d23e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sst_fit_sevi2=sst_fit_sevi.curvefit_coefficients\n",
    "phase_sevi=sst_fit_sevi.curvefit_coefficients.sel(param=\"b\")\n",
    "phase_sevi2=convert_phase(phase_sevi)\n",
    "print(phase_sevi2.min(),phase_sevi2.max())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bfc52de-efd7-4f1e-80fe-8def749289c7",
   "metadata": {},
   "source": [
    "### set dictionnaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d78ba16f-f230-4f27-9a92-88c2cf24c5dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "bounds=dict(a=(3,7),b=(0,30),c=(17,22))\n",
    "names=dict(a=\"Amplitude\",b=\"Phase\",c=\"Intercept\")\n",
    "cmap=dict(a=plt.cm.jet,b=plt.cm.gray,c=plt.cm.plasma)\n",
    "data_croco=dict(a=sst_croco.sel(param=\"a\"),b=phase_croco2,c=sst_croco.sel(param=\"c\"))\n",
    "data_sevi=dict(a=sst_fit_sevi2.sel(param=\"a\"),b=phase_sevi2,c=sst_fit_sevi2.sel(param=\"c\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ef291f5-edc0-4df5-a7ac-ca02e147ddbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_data(data,lon,lat,ax,name,model,**kw_plot):\n",
    "    kwargs_plot=dict(transform=ccrs.PlateCarree())\n",
    "    kwargs_plot.update(kw_plot)\n",
    "    ax.set_extent([-7,36,30,45],crs=ccrs.PlateCarree())\n",
    "    ax.coastlines()\n",
    "    ax.add_feature(cfeature.LAND, zorder=1, edgecolor='k')\n",
    "    ax.set_title(f'{name} SST {model}')\n",
    "    cf=ax.pcolormesh(lon,lat,data,**kwargs_plot)\n",
    "    cbar = fig.colorbar(cf, ax=ax, shrink=0.7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "259a6adb-8132-4c87-9681-c2aed7d55515",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97e8c9d2-af0b-424c-aa5e-2793e1ae64cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(nrows=3, ncols=2, figsize=(14, 12),subplot_kw=dict(projection=proj))\n",
    "for i,param in enumerate((\"a\",\"b\",\"c\")):\n",
    "   print(param)\n",
    "   data1=data_croco[param]\n",
    "   data1=data1.where(mask==1)  \n",
    "   data2=data_sevi[param]\n",
    "   vmin,vmax=bounds[param]\n",
    "   kw_plot=dict(vmin=vmin,vmax=vmax,cmap=cmap[param])\n",
    "   plot_data(data1,lon_croco,lat_croco,axes[i,0],names[param],\"CROCO\",**kw_plot)\n",
    "   plot_data(data2,lon_sevi,lat_sevi,axes[i,1],names[param],\"SEVIRI\",**kw_plot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dddeef1-af3c-435b-960d-189f0cd2a43d",
   "metadata": {},
   "outputs": [],
   "source": [
    "data=sst_mean.mean('time')\n",
    "data=data.where(mask==1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8b29976-f1a6-409e-8d71-59e8d43ac05f",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax=plt.subplots(1,1,figsize=(10,8),subplot_kw=dict(projection=proj))\n",
    "kwargs_plot=dict(transform=ccrs.PlateCarree(),cmap=plt.cm.jet,vmin=14,vmax=22)\n",
    "ax.set_extent([-7,36,30,45],crs=ccrs.PlateCarree())\n",
    "ax.coastlines()\n",
    "ax.add_feature(cfeature.LAND, zorder=1, edgecolor='k')\n",
    "ax.set_title(f'Amplitude SST CROCO')\n",
    "cf=ax.pcolormesh(lon,lat,data,**kwargs_plot)\n",
    "cbar = fig.colorbar(cf, ax=ax, shrink=0.7)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31f9a745-1567-4457-bad2-2be19dac7608",
   "metadata": {},
   "source": [
    "### check the results with scipy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8233c97-e5be-4a32-8d05-e83503bd54e0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "signal = sst2[:,400,600]\n",
    "ydata=signal\n",
    "x_0 = N.max(signal.data)-N.min(signal.data),0,N.mean(signal.data)\n",
    "#curve fit from scipy\n",
    "popt, pcov = curve_fit(model_curve, xdata, ydata)\n",
    "#least square fit\n",
    "x,cost = optimize_1freq(x_0,signal.data,loss=\"linear\")\n",
    "sig_fitted = model_1freq_trend(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "817129a8-7c8a-457e-9973-0ec224c3d399",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "#curve with curvefit from xarray\n",
    "\n",
    "fit= sst2.isel(Y=400,X=600).curvefit(\n",
    "       coords=\"month\", func=model_curve)#, p0=x_0)\n",
    "coeffs=fit.curvefit_coefficients.data.compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e549a800-b3b5-4338-bd11-baf0a7d656ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "coeffs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da4e80fb-e49e-474b-962e-d0e7584db39e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot\n",
    "fig,ax=plt.subplots(1,1,figsize=(10,8))\n",
    "ax.plot(xdata, ydata, 'b-', label='data')\n",
    "ax.plot(xdata,model_curve(xdata,*popt),'r-',label='scipy curve fit')\n",
    "ax.plot(xdata,model_curve(xdata,*coeffs),'g+',label='xarray wrapper curve fit')\n",
    "ax.plot(xdata,sig_fitted,'k*',markersize=2,label='least squares')\n",
    "plt.legend()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecb74e6d-b9f8-41a7-bb66-6ca660e0c952",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}