From f5a5a5cca1bc359af33efa8833c51c5ff421ad5e Mon Sep 17 00:00:00 2001
From: Leonardo Uieda <leo@uieda.com>
Date: Fri, 23 Aug 2024 12:02:39 -0300
Subject: [PATCH] Add NDVI calculation

Add a function that calculates NDVI for a given scene. It's better than
having to remember the equations all the time. The function can also
take care of adding some metadata and copying the metadata from the
original scene. The function allows customizing the names of the red and
NIR bands.
---
 doc/api/index.rst    | 25 ++++++++++++++++++++-----
 xlandsat/__init__.py |  1 +
 xlandsat/_indices.py | 41 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 62 insertions(+), 5 deletions(-)
 create mode 100644 xlandsat/_indices.py

diff --git a/doc/api/index.rst b/doc/api/index.rst
index 5a69dab..c8b2e3e 100644
--- a/doc/api/index.rst
+++ b/doc/api/index.rst
@@ -11,22 +11,37 @@ Input and output
 ----------------
 
 .. autosummary::
-   :toctree: generated/
+    :toctree: generated/
 
     load_scene
     save_scene
     load_panchromatic
 
-Processing
-----------
+Visualization
+-------------
 
 .. autosummary::
-   :toctree: generated/
+    :toctree: generated/
 
     composite
-    pansharpen
     equalize_histogram
     adjust_l1_colors
+
+Indices
+-------
+
+.. autosummary::
+    :toctree: generated/
+
+    ndvi
+
+Processing
+----------
+
+.. autosummary::
+    :toctree: generated/
+
+    pansharpen
     interpolate_missing
 
 Sample datasets
diff --git a/xlandsat/__init__.py b/xlandsat/__init__.py
index 0748713..f95519f 100644
--- a/xlandsat/__init__.py
+++ b/xlandsat/__init__.py
@@ -4,6 +4,7 @@
 from . import datasets
 from ._composite import composite
 from ._enhancement import adjust_l1_colors, equalize_histogram
+from ._indices import ndvi
 from ._interpolation import interpolate_missing
 from ._io import load_panchromatic, load_scene, save_scene
 from ._pansharpen import pansharpen
diff --git a/xlandsat/_indices.py b/xlandsat/_indices.py
new file mode 100644
index 0000000..f2237f8
--- /dev/null
+++ b/xlandsat/_indices.py
@@ -0,0 +1,41 @@
+# Copyright (c) 2022 The xlandsat developers.
+# Distributed under the terms of the MIT License.
+# SPDX-License-Identifier: MIT
+"""
+Calculate indices based on data in xarray.Datasets
+"""
+
+
+def ndvi(scene, red_band="red", nir_band="nir"):
+    """
+    Normalized difference vegetation index
+
+    Calculate the NDVI for the given scene, defined as:
+
+    .. math::
+
+        NDVI = \dfrac{NIR - Red}{NIR + Red}
+
+    Parameters
+    ----------
+    scene : :class:`xarray.Dataset`
+        A Landsat scene, as read with :func:`xlandsat.load_scene`.
+    red_band : str
+        The name of the variable in ``scene`` that corresponds to the red band.
+    nir_band : str
+        The name of the variable in ``scene`` that corresponds to the NIR band.
+
+    Returns
+    -------
+    ndvi : :class:`xarray.DataArray`
+        The calculated NDVI, with the metadata attributes from the original
+        scene.
+    """
+    red = scene[red_band]
+    nir = scene[nir_band]
+    result = (nir - red) / (nir + red)
+    result.name = "ndvi"
+    attrs = {"long_name": "normalized difference vegetation index"}
+    attrs.update(scene.attrs)
+    result = result.assign_attrs(attrs)
+    return result