.. _example-nonlinear-oscillator:

====================
Nonlinear Oscillator
====================

This example demonstrates building an ODE model using the **Nonlinear Oscillator** test case, which is included in the NeurEco installation package.

.. contents:: On this page
   :local:
   :depth: 2

Problem Description
===================

The Nonlinear Oscillator test case describes a Duffing oscillator governed by:

.. math::

   \ddot{x} + \delta \dot{x} + \alpha x + \beta x^3 = f(t)

where:

- :math:`\delta = 0.22` (damping)
- :math:`\alpha = 1` (linear stiffness)
- :math:`\beta = 875` (nonlinear stiffness)

**Input**: The forcing function :math:`f(t)`

**Output**: The displacement :math:`x(t)` (with added noise)

The parameters were chosen to produce strong nonlinearity.

Test Case Data
==============

The test case includes:

- **Training**: 2 trajectories (750 time steps each)
- **Validation**: 1 trajectory (1501 time steps)
- **Test**: 1 trajectory (1501 time steps)

Files are organized as:

.. code-block:: text

   data/
   |-- train/
   |   |-- train_exc_1.csv
   |   |-- train_out_1.csv
   |   |-- train_exc_2.csv
   |   `-- train_out_2.csv
   |-- valid/
   |   |-- valid_exc_1.csv
   |   `-- valid_out_1.csv
   `-- test/
       |-- test_exc_1.csv
       `-- test_out_1.csv

Setup
=====

1. Create a working directory (e.g., ``NonLinearOscillator``)
2. Extract the test case data into this directory
3. Create a Python script (e.g., ``nonlinear_oscillator.py``)

Complete Example
================

.. code-block:: python

   """
   NeurEco ODE Example: Nonlinear Oscillator

   This example uses the Nonlinear Oscillator test case included
   with the NeurEco installation.
   """

   from NeurEco import NeurEcoDynamic as Dynamic
   import numpy as np
   import matplotlib.pyplot as plt
   import os


   def load_ode_dataset(directory):
       """Load excitation and output arrays (time in column 0) from a directory."""
       excitations = []
       outputs = []

       for file_name in sorted(os.listdir(directory)):
           data = np.genfromtxt(
               os.path.join(directory, file_name),
               skip_header=True,
               delimiter=","
           )
           # CSV column 0 is time; keep it as column 0 of the arrays we pass to NeurEco
           if "exc" in file_name:
               excitations.append(data)   # column 0 = time, remaining columns = excitations
           elif "out" in file_name:
               outputs.append(data)       # column 0 = time, remaining columns = outputs

       return excitations, outputs


   # ===================
   # Step 1: Load Data
   # ===================

   train_excitations, train_outputs = load_ode_dataset("./data/train")
   valid_excitations, valid_outputs = load_ode_dataset("./data/valid")
   test_excitations, test_outputs = load_ode_dataset("./data/test")

   print(f"Training trajectories: {len(train_excitations)}")
   print(f"Validation trajectories: {len(valid_excitations)}")
   print(f"Test trajectories: {len(test_excitations)}")


   # ===================
   # Step 2: Build Model
   # ===================

   builder = Dynamic.ODE()

   # For the Duffing oscillator:
   # - Steady state: when f(t) = 0, x = 0 is a stable equilibrium
   # - Hidden states: 2nd order system suggests ~2-3 hidden states

   # Each array: column 0 = time, remaining columns = excitations / outputs.
   config = Dynamic.ODEBuildConfig(
       # Training data
       train_exc_list=train_excitations,
       train_out_list=train_outputs,
       # Validation data
       valid_exc_list=valid_excitations,
       valid_out_list=valid_outputs,
       # Steady state constraints
       steady_state_exc=np.array([0]),
       steady_state_out=np.array([0]),
       # Model architecture
       min_hidden_states=1,
       max_hidden_states=3,
       # Normalization
       inputs_normalize_per_feature=True,
       outputs_normalize_per_feature=True,
       # Output paths
       checkpoint_address="./NLOscillator/model.checkpoint",
       write_model_to="./NLOscillator/model.ernn",
   )

   status = builder.build(config)

   if status == 0:
       print("Build completed successfully!")


   # =====================
   # Step 3: Load and Inspect Model
   # =====================

   evaluator = Dynamic.ODE()
   evaluator.load("./NLOscillator/model.ernn")

   print(f"\nModel Information:")
   print(f"  Inputs: {evaluator.get_input_count()}")
   print(f"  Outputs: {evaluator.get_output_count()}")
   print(f"  Hidden states: {evaluator.get_state_count()}")
   print(f"  Initialization steps: {evaluator.get_number_of_init_tsteps()}")
   print(f"  Steady state inputs: {evaluator.get_steady_input()}")
   print(f"  Steady state outputs: {evaluator.get_steady_output()}")


   # =====================
   # Step 4: Evaluate Model
   # =====================

   # Evaluate on all datasets
   # Each excitation array has time in its first column; predictions are returned time-first.
   pred_train = evaluator.evaluate(train_excitations)
   pred_valid = evaluator.evaluate(valid_excitations)
   pred_test = evaluator.evaluate(test_excitations)

   # Compute errors
   for i, (pred, ref) in enumerate(zip(pred_train, train_outputs)):
       error = evaluator.compute_error(pred[:, 1:], ref[:, 1:])
       print(f"Training trajectory {i+1}: L2 error = {100*error:.4f}%")

   for i, (pred, ref) in enumerate(zip(pred_valid, valid_outputs)):
       error = evaluator.compute_error(pred[:, 1:], ref[:, 1:])
       print(f"Validation trajectory {i+1}: L2 error = {100*error:.4f}%")

   for i, (pred, ref) in enumerate(zip(pred_test, test_outputs)):
       error = evaluator.compute_error(pred[:, 1:], ref[:, 1:])
       print(f"Test trajectory {i+1}: L2 error = {100*error:.4f}%")


   # =====================
   # Step 5: Visualize Results
   # =====================

   fig, axes = plt.subplots(1, 2, figsize=(12, 4))

   # Plot validation results
   # Time is column 0 of each array; the displacement output is column 1
   ax = axes[0]
   ax.plot(valid_outputs[0][:, 0], valid_outputs[0][:, 1], 'r-', label="Reference")
   ax.plot(pred_valid[0][:, 0], pred_valid[0][:, 1], 'k--', label="Prediction")
   ax.set_xlabel("Time")
   ax.set_ylabel("Displacement x(t)")
   ax.set_title("Validation Trajectory")
   ax.legend()
   ax.grid(True)

   # Plot test results
   ax = axes[1]
   ax.plot(test_outputs[0][:, 0], test_outputs[0][:, 1], 'r-', label="Reference")
   ax.plot(pred_test[0][:, 0], pred_test[0][:, 1], 'k--', label="Prediction")
   ax.set_xlabel("Time")
   ax.set_ylabel("Displacement x(t)")
   ax.set_title("Test Trajectory")
   ax.legend()
   ax.grid(True)

   plt.suptitle("Nonlinear Oscillator: ODE Model Predictions")
   plt.tight_layout()
   plt.show()


   # =====================
   # Cleanup
   # =====================

   builder.delete()
   evaluator.delete()

Key Concepts Demonstrated
=========================

**Steady-State Constraints**

For the Duffing oscillator, when the forcing :math:`f(t) = 0`, the system has a stable equilibrium at :math:`x = 0`. Providing this steady-state information helps NeurEco build a more physically consistent model:

.. code-block:: python

   steady_state_exc=np.array([0]),
   steady_state_out=np.array([0]),

**Hidden State Bounds**

The governing equation is second-order, suggesting the system has two internal states (position and velocity). We allow up to 3 hidden states to account for the strong nonlinearity:

.. code-block:: python

   min_hidden_states=1,
   max_hidden_states=3,

**Loading Multiple Trajectories**

The training data contains two separate trajectories. These are passed as lists,
each array holding the time in its first column. The loader keeps every CSV file as
loaded (no time split) and sorts the file names, so that the i-th excitation array
belongs to the same trajectory as the i-th output array:

.. code-block:: python

   train_exc_list=train_excitations, # List of 2 excitation arrays (time in column 0)
   train_out_list=train_outputs,     # List of 2 output arrays (time in column 0)

Next Steps
==========

- Try different ``max_hidden_states`` values to see the trade-off between accuracy and model size
- Experiment with initialization during evaluation (see :ref:`ode-evaluating`)
- Export the model to FMU format (see :ref:`dynamic-exporting`)