diff --git a/.github/workflows/emscripten.yml b/.github/workflows/emscripten.yml
index 759c057..ea761c4 100644
--- a/.github/workflows/emscripten.yml
+++ b/.github/workflows/emscripten.yml
@@ -9,7 +9,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - name: Build WASM wheel
-        uses: pypa/cibuildwheel@v3.0.0
+        uses: pypa/cibuildwheel@v3.0.1
         env:
           CIBW_PLATFORM: pyodide
       - name: Upload package
diff --git a/.github/workflows/static.yml b/.github/workflows/static.yml
index 49765ec..f8b0fa4 100644
--- a/.github/workflows/static.yml
+++ b/.github/workflows/static.yml
@@ -9,7 +9,7 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-      - uses: prefix-dev/setup-pixi@v0.8.10
+      - uses: prefix-dev/setup-pixi@v0.8.11
         with:
           environments: static
           cache: true
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index bdd930c..61c518b 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -16,7 +16,7 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-      - uses: prefix-dev/setup-pixi@v0.8.10
+      - uses: prefix-dev/setup-pixi@v0.8.11
         with:
           environments: >-
             dev
diff --git a/.github/workflows/wheel.yml b/.github/workflows/wheel.yml
index 7a6a74c..90e9fdf 100644
--- a/.github/workflows/wheel.yml
+++ b/.github/workflows/wheel.yml
@@ -8,7 +8,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
-      - uses: prefix-dev/setup-pixi@v0.8.10
+      - uses: prefix-dev/setup-pixi@v0.8.11
         with:
           environments: dev
           cache: true
@@ -33,7 +33,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - name: Build wheels
-        uses: pypa/cibuildwheel@v3.0.0
+        uses: pypa/cibuildwheel@v3.0.1
         env:
           CIBW_SKIP: "*_i686 *_ppc64le *_s390x *_universal2 *-musllinux_*"
           CIBW_PROJECT_REQUIRES_PYTHON: ">=3.10"
diff --git a/.readthedocs.yml b/.readthedocs.yml
index 91913f3..eb05b26 100644
--- a/.readthedocs.yml
+++ b/.readthedocs.yml
@@ -3,9 +3,16 @@ build:
   os: ubuntu-22.04
   tools:
     python: "3.13"
+  apt_packages:
+    - default-jre # for plantuml
+    - graphviz # for plantuml
+  jobs:
+    pre_install:
+      - bash scripts/rtd_pre_install.sh
 
 sphinx:
    configuration: doc/conf.py
+   fail_on_warning: true
 
 python:
   install:
diff --git a/README.rst b/README.rst
index a14365f..1f27038 100644
--- a/README.rst
+++ b/README.rst
@@ -1,4 +1,4 @@
-FastCan: A Fast Canonical-Correlation-Based Greedy Search Algorithm
+fastcan: A fast canonical-correlation-based greedy search algorithm
 ===================================================================
 |conda| |Codecov| |CI| |Doc| |PythonVersion| |PyPi| |ruff| |pixi|
 
@@ -26,7 +26,7 @@ FastCan: A Fast Canonical-Correlation-Based Greedy Search Algorithm
 .. |pixi| image:: https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/prefix-dev/pixi/main/assets/badge/v0.json&style=flat-square
    :target: https://pixi.sh
 
-FastCan is a greedy search algorithm that supports:
+fastcan is a greedy search algorithm that supports:
 
 #. Feature selection
 
@@ -54,7 +54,7 @@ Check `Home Page <https://fastcan.readthedocs.io/en/latest/>`_ for more informat
 Installation
 ------------
 
-Install **FastCan** via PyPi:
+Install **fastcan** via PyPi:
 
 * Run ``pip install fastcan``
 
@@ -90,21 +90,21 @@ array([0.34617598, 0.95815008])
 
 NARX Time Series Modelling
 --------------------------
-FastCan can be used for system identification.
+fastcan can be used for system identification.
 In particular, we provide a submodule `fastcan.narx` to build Nonlinear AutoRegressive eXogenous (NARX) models.
 For more information, check our `Home Page <https://fastcan.readthedocs.io/en/latest/>`_.
 
 
 Support Free-Threaded Wheels
 ----------------------------
-FastCan has support for free-threaded (also known as nogil) CPython 3.13.
+fastcan has support for free-threaded (also known as nogil) CPython 3.13.
 For more information about free-threaded CPython, check `how to install a free-threaded CPython <https://py-free-threading.github.io/installing_cpython/>`_.
 
 Support WASM Wheels
 -------------------
-FastCan is compiled to WebAssembly (WASM) wheels using `pyodide <https://github.com/pyodide/pyodide>`_, and they are available on the assets of GitHub releases.
+fastcan is compiled to WebAssembly (WASM) wheels using `pyodide <https://github.com/pyodide/pyodide>`_, and they are available on the assets of GitHub releases.
 You can try it in a `REPL <https://pyodide.org/en/stable/console.html>`_ directly in a browser.
-The WASM wheels of FastCan can be installed by
+The WASM wheels of fastcan can be installed by
 
 >>> import micropip # doctest: +SKIP
 >>> await micropip.install('URL of the wasm wheel (end with _wasm32.whl)') # doctest: +SKIP
@@ -112,13 +112,13 @@ The WASM wheels of FastCan can be installed by
 📝 **Note:** Due to the Cross-Origin Resource Sharing (CORS) block in web browsers,
 you may need `Allow CORS: Access-Control-Allow-Origin Chrome extension <https://chrome.google.com/webstore/detail/allow-cors-access-control/lhobafahddgcelffkeicbaginigeejlf>`_.
 
-📝 **Note:** The nightly wasm wheel of FastCan's dependency (i.e. scikit-learn) can be found in `Scientific Python Nightly Wheels <https://pypi.anaconda.org/scientific-python-nightly-wheels/simple/>`_.
+📝 **Note:** The nightly wasm wheel of fastcan's dependency (i.e. scikit-learn) can be found in `Scientific Python Nightly Wheels <https://pypi.anaconda.org/scientific-python-nightly-wheels/simple/>`_.
 
 
 Citation
 --------
 
-FastCan is a Python implementation of the following papers.
+fastcan is a Python implementation of the following papers.
 
 If you use the `h-correlation` method in your work please cite the following reference:
 
diff --git a/doc/conf.py b/doc/conf.py
index bc47a1b..078549f 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -41,8 +41,10 @@
     # Link to other project's documentation (see mapping below)
     "sphinx.ext.intersphinx",
     "sphinx_gallery.gen_gallery",
-    # for dropdown directive
+    # For dropdown directive
     "sphinx_design",
+    # For PlantUML diagrams
+    "sphinxcontrib.plantuml",
 ]
 
 # List of patterns, relative to source directory, that match files and
diff --git a/doc/diagram.puml b/doc/diagram.puml
new file mode 100644
index 0000000..54924a3
--- /dev/null
+++ b/doc/diagram.puml
@@ -0,0 +1,73 @@
+@startuml fastcan
+skinparam backgroundColor transparent
+!include https://raw.githubusercontent.com/plantuml-stdlib/C4-PlantUML/master/C4_Component.puml
+!include <logos/numpy>
+!include <logos/python>
+!include <logos/c>
+
+
+AddContainerTag("module", $legendText="module")
+AddContainerTag("db", $legendText="input/output", $sprite="numpy", $bgColor="gray", $fontColor="white", $borderColor="gray")
+AddComponentTag("python", $legendText="Python code", $sprite="python")
+AddComponentTag("cython", $legendText="Cython code", $sprite="c", $bgColor="gold", $fontColor="brown", $borderColor="gold")
+UpdateContainerBoundaryStyle($type="module", $legendText="module boundary")
+
+
+Container_Boundary(fastcan, "fastcan", $descr="A library for fast feature engineering and data preprocessing"){
+    Component(cancorr_fast, "cancorr_fast", $tags="cython", $descr="Fast canonical correlation based forward search")
+    Component(FastCan, "FastCan", $tags="python", $descr="Feature selector")
+    Component(minibatch, "minibatch", $tags="python", $descr="Prunes samples in batch-wise")
+    Component(refine, "refine", $tags="python", $descr="Refines selection of FastCan")
+
+
+    Rel(FastCan, cancorr_fast, "Sends features to", "arrays")
+    Rel(minibatch, cancorr_fast, "Sends samples to", "arrays")
+    Rel(FastCan, refine, "Sends selected features to", "arrays")
+    Rel(refine, cancorr_fast, "Sends features to", "arrays")
+}
+
+Container_Boundary(narx, "narx", $descr="A submodule for NARX modelling"){
+    Component(make_narx, "make_narx", $tags="python", $descr="Builder for NARX model instances")
+    Component(narx_fast, "narx_fast", $tags="cython", $descr="Fast computation of gradient and prediction for NARX models")
+    Component(NARX, "NARX", $tags="python", $descr="NARX model")
+    Component(time_shift, "time_shift", $tags="python", $descr="Transforming time-series to time-shifted features")
+    Component(poly, "poly", $tags="python", $descr="Nonlinearises features with polynomial basis functions")
+    ' Component(tp2fd, "tp2fd", $tags="python", $descr="Converts time_shift ids and poly ids to feat ids and delay ids")
+    Component(print_narx, "print_narx", $tags="python", $descr="Prints NARX model summary")
+
+    Rel(NARX, print_narx, "Sends NARX model to", "NARX model")
+    ' Rel(make_narx, poly, "Makes polynomial features using", "unique id numbers")
+    Rel(make_narx, time_shift, "Sends time series to", "arrays")
+    ' Rel(make_narx, tp2fd, "Sends time_shift ids and poly ids to", "unique id numbers")
+    Rel(time_shift, poly, "Sends time-shifted features to", "arrays")
+    Rel(poly, FastCan, "Sends polynomial features to", "arrays")
+    ' Rel(tp2fd, NARX, "Sends feat ids and delay ids to", "unique id numbers")
+    Rel(NARX, narx_fast, "Sends initial conditions of inputs, prediction and gradients to", "fit, predict")
+
+}
+
+Person(person, "User", $descr="A data scientist or developer using NumPy, SciPy, and scikit-learn")
+ContainerDb(output, "fastcan output", $tags="db", $techn="indices", $descr="Selected indices of features or samples")
+ContainerDb(input, "fastcan input", $tags="db", $techn="arrays, allow multi-output", $descr="Input data")
+ContainerDb(narx_output, "narx output", $tags="db", $techn="arrays", $descr="Prediction and gradients of NARX model")
+ContainerDb(narx_input, "narx input", $tags="db", $techn="arrays, allow nan, allow multi-output", $descr="Time-series data")
+
+Rel(input, FastCan, "Sends features to", "arrays")
+Rel(input, minibatch, "Sends samples to", "arrays")
+Rel(narx_input, make_narx, "Sends time series to", "arrays")
+Rel(narx_input, NARX, "Sends time series to", "arrays")
+
+Rel(cancorr_fast, output, "Sends selected indices to", "indices")
+Rel(output, refine, "Sends selected indices to", "indices")
+Rel(narx_fast, narx_output, "Sends prediction and gradients to", "arrays")
+Rel(output, NARX, "Sends selected polynomial features to", "indices")
+
+Rel(person, input, "Processes arrays using", "NumPy, scikit-learn pipeline")
+Rel(person, narx_input, "Processes time series using", "NumPy, scikit-learn pipeline")
+
+Container(optimizer, "SciPy Optimiser", "module", $tags="module", $descr="Minimises objective functions using prediction errors and gradients")
+Rel(narx_output, optimizer, "Sends prediction and gradients to", "arrays")
+Rel(optimizer, NARX, "Updates coefficients for", "arrays")
+
+SHOW_LEGEND()
+@enduml
\ No newline at end of file
diff --git a/doc/index.rst b/doc/index.rst
index 071cafd..06fde70 100644
--- a/doc/index.rst
+++ b/doc/index.rst
@@ -9,7 +9,11 @@
 
 .. include:: ../README.rst
 
+Architecture Diagram
+--------------------
 
+.. uml:: diagram.puml
+   :align: center
 
 API Reference
 -------------
@@ -49,7 +53,7 @@ Useful Links
 API Compatibility
 -----------------
 
-The API of this package is align with scikit-learn.
+The API of this library is align with scikit-learn.
 
 .. |sklearn| image:: images/scikit-learn-logo-notext.png
    :width: 100pt
diff --git a/pixi.lock b/pixi.lock
index cd85c91..29cd125 100644
--- a/pixi.lock
+++ b/pixi.lock
@@ -562,6 +562,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/3c/dd/018ce05c532a22007ac58d4f45232514cd9d6dd0ee1dc374e309db830983/sphinx_basic_ng-1.0.0b2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/c6/43/65c0acbd8cc6f50195a3a1fc195c404988b15c67090e73c7a41a9f57d6bd/sphinx_design-0.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/77/c7/52b48aec16b26c52aba854d03a3a31e0681150301dac1bea2243645a69e7/sphinx_gallery-0.19.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/4d/e0/91ee50f1a020e2ed48d370a054f94b012ba0d757214a420ac43c9327f818/sphinxcontrib-plantuml-0.30.tar.gz
       - pypi: https://files.pythonhosted.org/packages/5d/85/9ebeae2f76e9e77b952f4b274c27238156eae7979c5421fba91a28f4970d/sphinxcontrib_applehelp-2.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/35/7a/987e583882f985fe4d7323774889ec58049171828b58c2217e7f79cdf44e/sphinxcontrib_devhelp-2.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/0a/7b/18a8c0bcec9182c05a0b3ec2a776bba4ead82750a55ff798e8d406dae604/sphinxcontrib_htmlhelp-2.1.0-py3-none-any.whl
@@ -840,6 +841,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/3c/dd/018ce05c532a22007ac58d4f45232514cd9d6dd0ee1dc374e309db830983/sphinx_basic_ng-1.0.0b2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/c6/43/65c0acbd8cc6f50195a3a1fc195c404988b15c67090e73c7a41a9f57d6bd/sphinx_design-0.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/77/c7/52b48aec16b26c52aba854d03a3a31e0681150301dac1bea2243645a69e7/sphinx_gallery-0.19.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/4d/e0/91ee50f1a020e2ed48d370a054f94b012ba0d757214a420ac43c9327f818/sphinxcontrib-plantuml-0.30.tar.gz
       - pypi: https://files.pythonhosted.org/packages/5d/85/9ebeae2f76e9e77b952f4b274c27238156eae7979c5421fba91a28f4970d/sphinxcontrib_applehelp-2.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/35/7a/987e583882f985fe4d7323774889ec58049171828b58c2217e7f79cdf44e/sphinxcontrib_devhelp-2.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/0a/7b/18a8c0bcec9182c05a0b3ec2a776bba4ead82750a55ff798e8d406dae604/sphinxcontrib_htmlhelp-2.1.0-py3-none-any.whl
@@ -1090,6 +1092,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/3c/dd/018ce05c532a22007ac58d4f45232514cd9d6dd0ee1dc374e309db830983/sphinx_basic_ng-1.0.0b2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/c6/43/65c0acbd8cc6f50195a3a1fc195c404988b15c67090e73c7a41a9f57d6bd/sphinx_design-0.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/77/c7/52b48aec16b26c52aba854d03a3a31e0681150301dac1bea2243645a69e7/sphinx_gallery-0.19.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/4d/e0/91ee50f1a020e2ed48d370a054f94b012ba0d757214a420ac43c9327f818/sphinxcontrib-plantuml-0.30.tar.gz
       - pypi: https://files.pythonhosted.org/packages/5d/85/9ebeae2f76e9e77b952f4b274c27238156eae7979c5421fba91a28f4970d/sphinxcontrib_applehelp-2.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/35/7a/987e583882f985fe4d7323774889ec58049171828b58c2217e7f79cdf44e/sphinxcontrib_devhelp-2.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/0a/7b/18a8c0bcec9182c05a0b3ec2a776bba4ead82750a55ff798e8d406dae604/sphinxcontrib_htmlhelp-2.1.0-py3-none-any.whl
@@ -1340,6 +1343,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/3c/dd/018ce05c532a22007ac58d4f45232514cd9d6dd0ee1dc374e309db830983/sphinx_basic_ng-1.0.0b2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/c6/43/65c0acbd8cc6f50195a3a1fc195c404988b15c67090e73c7a41a9f57d6bd/sphinx_design-0.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/77/c7/52b48aec16b26c52aba854d03a3a31e0681150301dac1bea2243645a69e7/sphinx_gallery-0.19.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/4d/e0/91ee50f1a020e2ed48d370a054f94b012ba0d757214a420ac43c9327f818/sphinxcontrib-plantuml-0.30.tar.gz
       - pypi: https://files.pythonhosted.org/packages/5d/85/9ebeae2f76e9e77b952f4b274c27238156eae7979c5421fba91a28f4970d/sphinxcontrib_applehelp-2.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/35/7a/987e583882f985fe4d7323774889ec58049171828b58c2217e7f79cdf44e/sphinxcontrib_devhelp-2.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/0a/7b/18a8c0bcec9182c05a0b3ec2a776bba4ead82750a55ff798e8d406dae604/sphinxcontrib_htmlhelp-2.1.0-py3-none-any.whl
@@ -1580,6 +1584,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/3c/dd/018ce05c532a22007ac58d4f45232514cd9d6dd0ee1dc374e309db830983/sphinx_basic_ng-1.0.0b2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/c6/43/65c0acbd8cc6f50195a3a1fc195c404988b15c67090e73c7a41a9f57d6bd/sphinx_design-0.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/77/c7/52b48aec16b26c52aba854d03a3a31e0681150301dac1bea2243645a69e7/sphinx_gallery-0.19.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/4d/e0/91ee50f1a020e2ed48d370a054f94b012ba0d757214a420ac43c9327f818/sphinxcontrib-plantuml-0.30.tar.gz
       - pypi: https://files.pythonhosted.org/packages/5d/85/9ebeae2f76e9e77b952f4b274c27238156eae7979c5421fba91a28f4970d/sphinxcontrib_applehelp-2.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/35/7a/987e583882f985fe4d7323774889ec58049171828b58c2217e7f79cdf44e/sphinxcontrib_devhelp-2.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/0a/7b/18a8c0bcec9182c05a0b3ec2a776bba4ead82750a55ff798e8d406dae604/sphinxcontrib_htmlhelp-2.1.0-py3-none-any.whl
@@ -3605,13 +3610,14 @@ packages:
 - pypi: ./
   name: fastcan
   version: 0.4.0
-  sha256: 327259867342e372de54fcbb50da07c104703821040da6cbff3f125c846a1374
+  sha256: d6ce8adf0fa75fd949496cb6cb7602278dae2a1029b5ff4803901e07ff4a7b47
   requires_dist:
   - scikit-learn>=1.7.0
   - furo ; extra == 'docs'
   - matplotlib ; extra == 'docs'
   - sphinx-gallery ; extra == 'docs'
   - sphinx-design ; extra == 'docs'
+  - sphinxcontrib-plantuml ; extra == 'docs'
   requires_python: '>=3.10'
   editable: true
 - conda: https://conda.anaconda.org/conda-forge/noarch/font-ttf-dejavu-sans-mono-2.37-hab24e00_0.tar.bz2
@@ -11639,6 +11645,15 @@ packages:
   - flake8 ; extra == 'test'
   - mypy ; extra == 'test'
   requires_python: '>=3.5'
+- pypi: https://files.pythonhosted.org/packages/4d/e0/91ee50f1a020e2ed48d370a054f94b012ba0d757214a420ac43c9327f818/sphinxcontrib-plantuml-0.30.tar.gz
+  name: sphinxcontrib-plantuml
+  version: '0.30'
+  sha256: 2a1266ca43bddf44640ae44107003df4490de2b3c3154a0d627cfb63e9a169bf
+  requires_dist:
+  - sphinx>=1.6
+  - pytest ; extra == 'test'
+  - pillow ; extra == 'test'
+  - flake8 ; extra == 'test'
 - pypi: https://files.pythonhosted.org/packages/27/83/859ecdd180cacc13b1f7e857abf8582a64552ea7a061057a6c716e790fce/sphinxcontrib_qthelp-2.0.0-py3-none-any.whl
   name: sphinxcontrib-qthelp
   version: 2.0.0
diff --git a/pyproject.toml b/pyproject.toml
index 36b2cf3..faffea4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -33,7 +33,7 @@ download = "https://pypi.org/project/fastcan/#files"
 tracker = "https://github.com/scikit-learn-contrib/fastcan/issues"
 
 [project.optional-dependencies]
-docs = ["furo", "matplotlib", "sphinx_gallery", "sphinx-design"]
+docs = ["furo", "matplotlib", "sphinx_gallery", "sphinx-design", "sphinxcontrib-plantuml"]
 
 [build-system]
 requires = ["meson-python>=0.18.0", "Cython>=3.1.0", "scikit-learn>=1.7.0"]
diff --git a/scripts/rtd_pre_install.sh b/scripts/rtd_pre_install.sh
new file mode 100644
index 0000000..78fe7cc
--- /dev/null
+++ b/scripts/rtd_pre_install.sh
@@ -0,0 +1,18 @@
+#!/bin/bash
+
+# Stop and exit on error
+set -euox pipefail
+
+# Check for required tools
+java -version
+dot -V
+
+# This folder is on PATH and does not require sudo
+# Download latest plantuml.jar from github
+curl -o ${READTHEDOCS_VIRTUALENV_PATH}/bin/plantuml.jar -L https://github.com/plantuml/plantuml/releases/latest/download/plantuml.jar
+# Create an executable script for plantuml
+printf '#!/bin/bash\nexec java -Djava.awt.headless=true -jar ${READTHEDOCS_VIRTUALENV_PATH}/bin/plantuml.jar "$@"' > ${READTHEDOCS_VIRTUALENV_PATH}/bin/plantuml
+chmod +x ${READTHEDOCS_VIRTUALENV_PATH}/bin/plantuml
+
+# Check plantuml version
+plantuml -version
diff --git a/tests/test_narx.py b/tests/test_narx.py
index c72af2d..f8f1b03 100644
--- a/tests/test_narx.py
+++ b/tests/test_narx.py
@@ -614,4 +614,4 @@ def test_auto_reg():
     )
     model.fit(X, y)
     y_pred = model.predict(X, y_init=y[: model.max_delay_])
-    assert r2_score(y, model.predict(X, y_init=y)) > 0.5
+    assert r2_score(y, y_pred) > 0.5