Repository: aeon-toolkit/aeon
Pull request #1418 (2024-04-12 18:10:26)

[DOC] Correct `WeightedEnsembleClassifier` parameter docstring
<!--
Thanks for contributing a pull request! Please ensure you have taken a look at our
contribution guide: https://www.aeon-toolkit.org/en/stable/contributing.html

Feel free to delete sections of this template if they do not apply to your PR,
avoid submitting a blank template or empty sections.
-->

#### What does this implement/fix? Explain your changes.

<!--
A clear and concise description of what you have implemented.
-->

The current docstring for the `classifiers` parameter of `WeightedEnsembleClassifier` is not correct.

This change updates the docstring to match what other ensemble docstrings look like (see [forecasing.compose._ensemble.py](https://github.com/aeon-toolkit/aeon/blob/5f8974f3fca240f66cb52a626658931489871962/aeon/forecasting/compose/_ensemble.py#L46-L47))

#### Does your contribution introduce a new dependency? If yes, which one?

No

### PR checklist

<!--
Please go through the checklist below. Please feel free to remove points if they are not applicable.
-->

##### For all contributions
- [x] I've added myself to the [list of contributors](https://github.com/aeon-toolkit/aeon/blob/main/.all-contributorsrc). Alternatively, you can use the [@all-contributors](https://allcontributors.org/docs/en/bot/usage) bot to do this for you.
- [x] The PR title starts with either [ENH], [MNT], [DOC], [BUG], [REF], [DEP] or [GOV] indicating whether the PR topic is related to enhancement, maintenance, documentation, bugs, refactoring, deprecation or governance.

<!--
Thanks for contributing!
-->


aeon-actions-bot[bot]: 
## Thank you for contributing to `aeon`

I have added the following labels to this PR based on the title: **[ $\\color{#F3B9F8}{\\textsf{documentation}}$ ]**.
I have added the following labels to this PR based on the changes made: **[ $\\color{#BCAE15}{\\textsf{classification}}$ ]**. Feel free to change these if they do not properly represent the PR.

The [Checks](https://github.com/aeon-toolkit/aeon/pull/1418/checks) tab will show the status of our automated tests. You can click on individual test runs in the tab or "Details" in the panel below to see more information if there is a failure.

If our `pre-commit` code quality check fails, any trivial fixes will automatically be pushed to your PR unless it is a draft.

Don't hesitate to ask questions on the `aeon` [Slack](
https://join.slack.com/t/aeon-toolkit/shared_invite/zt-22vwvut29-HDpCu~7VBUozyfL_8j3dLA) channel if you have any.
    

Patch:
diff --git a/.github/workflows/github_maintenace.yml b/.github/workflows/github_maintenace.yml
index 3a483385..c50ff2f6 100644
--- a/.github/workflows/github_maintenace.yml
+++ b/.github/workflows/github_maintenace.yml
@@ -23,13 +23,13 @@ jobs:
           private-key: ${{ secrets.PR_APP_KEY }}
 
       - name: Stale Branches
-        uses: crs-k/stale-branches@v4.1.0
+        uses: crs-k/stale-branches@v5.0.0
         with:
           repo-token: ${{ steps.app-token.outputs.token }}
           days-before-stale: 140
           days-before-delete: 9999  # todo take down to 175
           comment-updates: true
           tag-committer: true
-          stale-branch-label: "stale"
+          stale-branch-label: "stale branch"
           compare-branches: 'info'
           pr-check: true
diff --git a/.github/workflows/pr_precommit.yml b/.github/workflows/pr_precommit.yml
index 882c8bb7..062e42aa 100644
--- a/.github/workflows/pr_precommit.yml
+++ b/.github/workflows/pr_precommit.yml
@@ -26,7 +26,7 @@ jobs:
           python-version: "3.10"
 
       - name: Get changed files
-        uses: tj-actions/changed-files@v44.0.0
+        uses: tj-actions/changed-files@v44.0.1
         id: changed-files
 
       - name: List changed files
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index e560c164..b28445dd 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -29,7 +29,7 @@ repos:
         args: [ "--create", "--python-folders", "aeon" ]
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.3.5
+    rev: v0.3.7
     hooks:
       - id: ruff
         args: [ "--fix"]
@@ -55,7 +55,7 @@ repos:
         args: [ "--max-line-length=88", "--extend-ignore=E203" ]
 
   - repo: https://github.com/psf/black
-    rev: 24.3.0
+    rev: 24.4.0
     hooks:
       - id: black
         language_version: python3
diff --git a/aeon/base/_base_series.py b/aeon/base/_base_series.py
index 79433a7b..9a3f94e7 100644
--- a/aeon/base/_base_series.py
+++ b/aeon/base/_base_series.py
@@ -1,4 +1,11 @@
-"""Base class for estimators that fit single (possibly multivariate) time series."""
+"""Base class for estimators that fit single time series.
+
+This time series can be univariate or multivariate. The time series can potentially
+contain missing values.
+"""
+
+__maintainer__ = ["TonyBagnall", "MatthewMiddlehurst"]
+__all__ = ["BaseSeriesEstimator"]
 
 import numpy as np
 import pandas as pd
@@ -18,40 +25,33 @@ VALID_INPUT_TYPES = [pd.DataFrame, pd.Series, np.ndarray]
 class BaseSeriesEstimator(BaseEstimator):
     """Base class for estimators that use single (possibly multivariate) time series.
 
-    Provides functions that are common to BaseSeriesEstimator objects, including
-    BaseSeriesTransformer and BaseSegmenter, for the checking and
-    conversion of input to fit, predict and transform, where relevant.
+    Provides functions that are common to BaseSeriesEstimator objects for the checking
+    and conversion of input to fit, predict and transform, where relevant.
 
-    It also stores the common default tags used by all the subclasses and meta data
-    describing the characteristics of time series passed to ``fit``.
+    It also stores the common default tags used by all the subclasses and metadata
+    describing the characteristics of time series passed to ``fit`` (or another method
+    if fit does not exist).
 
-    input and internal data format
+    Input and internal data format (where m is the number of time points and d is the
+    number of channels):
         Univariate series:
-            Numpy array:
-            shape `(m,)`, `(m, 1)` or `(1, m)`. if ``self`` has no multivariate
-            capability, i.e.``self.get_tag(
-            ""capability:multivariate") == False``, all are converted to 1D
-            numpy `(m,)`
-            if ``self`` has multivariate capability, converted to 2D numpy `(m,1)` or
-            `(1, m)` depending on axis
-            pandas DataFrame or Series:
-            DataFrame single column shape `(m,1)`, `(1,m)` or Series shape `(m,)`
-            if ``self`` has no multivariate capability, all converted to Series `(m,)`
-            if ``self`` has multivariate capability, all converted to Pandas DataFrame
-            shape `(m,1)`, `(1,m)` depending on axis
-
+            np.ndarray, shape ``(m,)``, ``(m, 1)`` or ``(1, m)`` depending on axis.
+            This is converted to a 2D np.ndarray internally.
+            pd.DataFrame, shape ``(m, 1)`` or ``(1, m)`` depending on axis.
+            pd.Series, shape ``(m,)``.
         Multivariate series:
-            Numpy array, shape `(m,d)` or `(d,m)`.
-            pandas DataFrame `(m,d)` or `(d,m)`
+            np.ndarray array, shape ``(m, d)`` or ``(d, m)`` depending on axis.
+            pd.DataFrame ``(m, d)`` or ``(d, m)`` depending on axis.
 
     Parameters
     ----------
-    axis : int, default = 0
-        Axis along which to segment if passed a multivariate series (2D input). If axis
-        is 0, it is assumed each column is a time series and each row is a
-        timepoint. i.e. the shape of the data is ``(n_timepoints,n_channels)``.
-        ``axis == 1`` indicates the time series are in rows, i.e. the shape of the data
-        is ``(n_channels, n_timepoints)``.
+    axis : int
+        The time point axis of the input series if it is 2D. If ``axis==0``, it is
+        assumed each column is a time series and each row is a time point. i.e. the
+        shape of the data is ``(n_timepoints, n_channels)``. ``axis==1`` indicates
+        the time series are in rows, i.e. the shape of the data is
+        ``(n_channels, n_timepoints)``.
+        Setting this class variable will convert the input data to the chosen axis.
     """
 
     _tags = {
@@ -61,153 +61,183 @@ class BaseSeriesEstimator(BaseEstimator):
         "X_inner_type": "np.ndarray",  # one of VALID_INNER_TYPES
     }
 
-    def __init__(self, axis=0):
+    def __init__(self, axis):
         self.axis = axis
-        self.metadata_ = {}  # metadata/properties of data seen in fit
+        self.metadata_ = {}  # metadata/properties of data seen in fit/predict/transform
+
         super().__init__()
         _check_estimator_deps(self)
 
-    def _check_X(self, X):
-        """Check classifier input X is valid.
+    def _check_X(self, X, axis):
+        """Check input X is valid.
 
         Check if the input data is a compatible type, and that this estimator is
         able to handle the data characteristics. This is done by matching the
         capabilities of the estimator against the metadata for X for
-        univariate/multivariate, equal length/unequal length and no missing
-        values/missing values.
+        univariate/multivariate and no missing values/missing values.
 
         Parameters
         ----------
-        X : data structure
-           A valid aeon collection data structure. See
-           aeon.registry.COLLECTIONS_DATA_TYPES for details
-           on aeon supported data structures.
+        X: one of aeon.base._base_series.VALID_INPUT_TYPES
+            A valid aeon time series data structure. See
+            aeon.base._base_series.VALID_INPUT_TYPES for aeon supported types.
+        axis: int
+            The time point axis of the input series if it is 2D. If ``axis==0``, it is
+            assumed each column is a time series and each row is a time point. i.e. the
+            shape of the data is ``(n_timepoints,n_channels)``. ``axis==1`` indicates
+            the time series are in rows, i.e. the shape of the data is
+            ``(n_channels,n_timepoints)``.
 
         Returns
         -------
-        dict
-            Meta data about X, with flags:
-            metadata["missing_values"] : whether X has missing values or not
-            metadata["multivariate"] : whether X has more than one channel or not
-
-        See Also
-        --------
-        _convert_X : function that converts X after it has been checked.
+        metadata: dict
+            Metadata about X, with flags:
+            metadata["multivariate"]: whether X has more than one channel or not
+            metadata["n_channels"]: number of channels in X
+            metadata["missing_values"]: whether X has missing values or not
         """
-        # Checks: check valid type and axis
-        if type(X) not in VALID_INPUT_TYPES:
-            raise ValueError(
-                f"Error in input type should be one of "
-                f" {VALID_INNER_TYPES}, saw {type(X)}"
-            )
+        if axis > 1 or axis < 0:
+            raise ValueError(f"Input axis should be 0 or 1, saw {axis}")
+
+        # Checks: check valid dtype
         if isinstance(X, np.ndarray):
-            # Check valid shape
-            if X.ndim > 2:
-                raise ValueError("Should be 1D or 2D")
             if not (
                 issubclass(X.dtype.type, np.integer)
                 or issubclass(X.dtype.type, np.floating)
             ):
-                raise ValueError("np.ndarray must contain floats or ints")
+                raise ValueError("dtype for np.ndarray must be float or int")
         elif isinstance(X, pd.Series):
             if not pd.api.types.is_numeric_dtype(X):
-                raise ValueError("pd.Series must be numeric")
-        else:
+                raise ValueError("pd.Series dtype must be numeric")
+        elif isinstance(X, pd.DataFrame):
             if not all(pd.api.types.is_numeric_dtype(X[col]) for col in X.columns):
-                raise ValueError("pd.DataFrame must be numeric")
-        # If X is a single series dataframe, we squeeze it into Series in convert_X
-        X = X.squeeze()
+                raise ValueError("pd.DataFrame dtype must be numeric")
+        else:
+            raise ValueError(
+                f"Input type of X should be one of {VALID_INNER_TYPES}, saw {type(X)}"
+            )
+
         metadata = {}
-        metadata["multivariate"] = False
-        # Need to differentiate because a 1D series stored in a dataframe will have
-        # ndim=2. This case is dealt with in convert through squeezing to 1D
-        if X.ndim > 1:
+
+        # check if multivariate
+        channel_idx = 0 if axis == 1 else 1
+        if X.ndim > 2:
+            raise ValueError(
+                "X must have at most 2 dimensions for multivariate data, optionally 1 "
+                f"for univarate data. Found {X.ndim} dimensions"
+            )
+        elif X.ndim > 1 and X.shape[channel_idx] > 1:
             metadata["multivariate"] = True
+        else:
+            metadata["multivariate"] = False
+
+        metadata["n_channels"] = X.shape[channel_idx] if X.ndim > 1 else 1
+
+        # check if has missing values
         if isinstance(X, np.ndarray):
             metadata["missing_values"] = np.isnan(X).any()
         elif isinstance(X, pd.Series):
             metadata["missing_values"] = X.isna().any()
-        elif isinstance(X, pd.DataFrame):
+        else:
             metadata["missing_values"] = X.isna().any().any()
+
         allow_multivariate = self.get_tag("capability:multivariate")
         allow_univariate = self.get_tag("capability:univariate")
         allow_missing = self.get_tag("capability:missing_values")
         if metadata["missing_values"] and not allow_missing:
-            raise ValueError("Missing values not supported")
+            raise ValueError(
+                f"Missing values not supported by {self.__class__.__name__}"
+            )
         if metadata["multivariate"] and not allow_multivariate:
-            raise ValueError("Multivariate data not supported")
+            raise ValueError(
+                f"Multivariate data not supported by {self.__class__.__name__}"
+            )
         if not metadata["multivariate"] and not allow_univariate:
-            raise ValueError("Univariate data not supported")
+            raise ValueError(
+                f"Univariate data not supported by {self.__class__.__name__}"
+            )
+
         return metadata
 
-    def _check_y(self, y: VALID_INPUT_TYPES):
-        """Check y specific to segmentation.
+    def _convert_X(self, X, axis):
+        """Convert input X to internal estimator datatype.
+
+        Converts input X to the internal data type of the estimator using
+        self.get_tag("X_inner_type"). 1D numpy arrays are converted to 2D,
+        and the data will be transposed if the input axis does not match that of the
+        estimator.
+
+        Attempting to convert to a pd.Series for multivariate data or estimators will
+        raise an error.
 
-        y must be a univariate series
+        Parameters
+        ----------
+        X: one of aeon.base._base_series.VALID_INPUT_TYPES
+            A valid aeon time series data structure. See
+            aeon.base._base_series.VALID_INPUT_TYPES for aeon supported types.
+        axis: int
+            The time point axis of the input series if it is 2D. If ``axis==0``, it is
+            assumed each column is a time series and each row is a time point. i.e. the
+            shape of the data is ``(n_timepoints, n_channels)``. ``axis==1`` indicates
+            the time series are in rows, i.e. the shape of the data is
+            ``(n_channels, n_timepoints)``.
+
+        Returns
+        -------
+        X: one of aeon.base._base_series.VALID_INPUT_TYPES
+            Input time series with data structure of type self.get_tag("X_inner_type").
         """
-        if type(y) not in VALID_INPUT_TYPES:
-            raise ValueError(
-                f"Error in input type for y: it should be one of "
-                f"{VALID_INPUT_TYPES}, saw {type(y)}"
-            )
-        if isinstance(y, np.ndarray):
-            # Check valid shape
-            if y.ndim > 1:
-                raise ValueError(
-                    "Error in input type for y: y input as np.ndarray " "should be 1D"
-                )
-            if not (
-                issubclass(y.dtype.type, np.integer)
-                or issubclass(y.dtype.type, np.floating)
-            ):
-                raise ValueError(
-                    "Error in input type for y: y input must contain " "floats or ints"
-                )
-        elif isinstance(y, pd.Series):
-            if not pd.api.types.is_numeric_dtype(y):
-                raise ValueError(
-                    "Error in input type for y: y input as pd.Series must be numeric"
-                )
-        else:  # pd.DataFrame
-            if y.shape[1] > 2:
-                raise ValueError(
-                    "Error in input type for y: y input as pd.DataFrame "
-                    "should have a single "
-                    "column series"
-                )
+        if axis > 1 or axis < 0:
+            raise ValueError(f"Input axis should be 0 or 1, saw {axis}")
 
-            if not all(pd.api.types.is_numeric_dtype(y[col]) for col in y.columns):
-                raise ValueError(
-                    "Error in input type for y: y input as pd.DataFrame "
-                    "must be numeric"
-                )
+        inner_type = self.get_tag("X_inner_type")
+        if not isinstance(inner_type, list):
+            inner_type = [inner_type]
+        inner_names = [i.split(".")[-1] for i in inner_type]
 
-    def _convert_X(self, X, axis):
-        inner = self.get_tag("X_inner_type").split(".")[-1]
         input = type(X).__name__
-        if inner != input:
-            if inner == "ndarray":
+        if input not in inner_names:
+            if inner_names[0] == "ndarray":
                 X = X.to_numpy()
-            elif inner == "Series":
-                if input == "ndarray":
-                    X = pd.Series(X)
-            elif inner == "DataFrame":
+            elif inner_names[0] == "Series":
+                if self.get_tag("capability:multivariate"):
+                    raise ValueError(
+                        "Cannot convert to pd.Series for multivariate capable "
+                        "estimators"
+                    )
+                if X.ndim > 1:
+                    n_channels = X.shape[0] if axis == 1 else X.shape[1]
+                    if n_channels > 1:
+                        raise ValueError(
+                            "Cannot convert to pd.Series for multivariate data. Found "
+                            f"{n_channels} channels"
+                        )
+
+                X = X.squeeze()
+                X = pd.Series(X)
+            elif inner_names[0] == "DataFrame":
+                # converting a 1d array will create a 2d array in axis 0 format
+                transpose = False
+                if X.ndim == 1 and axis == 1:
+                    transpose = True
+
                 X = pd.DataFrame(X)
+
+                if transpose:
+                    X = X.T
             else:
-                tag = self.get_tag("X_inner_type")
-                raise ValueError(f"Unsupported inner type {inner} derived from {tag}")
-        if axis > 1 or axis < 0:
-            raise ValueError("Axis should be 0 or 1")
-        if not self.get_tag("capability:multivariate"):
-            X = X.squeeze()
-        elif X.ndim == 1:  # np.ndarray case make 2D
-            X = X.reshape(1, -1)
-        if X.ndim > 1:
-            if self.axis != axis:
-                X = X.T
+                raise ValueError(
+                    f"Unsupported inner type {inner_names[0]} derived from {inner_type}"
+                )
+
+        if X.ndim > 1 and self.axis != axis:
+            X = X.T
+        elif X.ndim == 1 and isinstance(X, np.ndarray):
+            X = X[np.newaxis, :] if self.axis == 1 else X[:, np.newaxis]
+
         return X
 
-    def _preprocess_series(self, X, axis=None):
+    def _preprocess_series(self, X, axis, store_metadata):
         """Preprocess input X prior to call to fit.
 
         Checks the characteristics of X, store metadata, checks self can handle
@@ -215,49 +245,64 @@ class BaseSeriesEstimator(BaseEstimator):
 
         Parameters
         ----------
-        X : one of VALID_INNER_TYPES
+        X: one of aeon.base._base_series.VALID_INPUT_TYPES
+            A valid aeon time series data structure. See
+            aeon.base._base_series.VALID_INPUT_TYPES for aeon supported types.
         axis: int or None
+            The time point axis of the input series if it is 2D. If ``axis==0``, it is
+            assumed each column is a time series and each row is a time point. i.e. the
+            shape of the data is ``(n_timepoints, n_channels)``. ``axis==1`` indicates
+            the time series are in rows, i.e. the shape of the data is
+            ``(n_channels, n_timepoints)``.
+            If None, the default class axis is used.
+        store_metadata: bool
+            If True, overwrite metadata with the new metadata from X.
 
         Returns
         -------
-        Data structure of type self.tags["X_inner_type"]
-
-        See Also
-        --------
-        _check_X : function that checks X is valid before conversion.
-        _convert_X : function that converts to inner type.
-        pass
+        X: one of aeon.base._base_series.VALID_INPUT_TYPES
+            Input time series with data structure of type self.get_tag("X_inner_type").
         """
         if axis is None:
             axis = self.axis
-        meta = self._check_X(X)
-        if len(self.metadata_) == 0:
+
+        meta = self._check_X(X, axis)
+        if store_metadata:
             self.metadata_ = meta
+
         return self._convert_X(X, axis)
 
-    def _postprocess_series(self, X, axis=None):
-        """Postprocess data X to revert to original shape."""
-        # If a univariate only transformer, return a univariate series
-        if not self.get_tag("capability:multivariate"):
-            return X.squeeze()
-        # If passed an axis, return with that axis
-        if axis is None or axis == self.axis:
-            return X
-        return X.T
-
-    @classmethod
-    def get_test_params(cls, parameter_set="default"):
-        """
-        Return testing parameter settings for the estimator.
+    def _postprocess_series(self, Xt, axis):
+        """Postprocess data Xt to revert to original shape.
 
         Parameters
         ----------
-        parameter_set : str, default="default"
+        Xt: one of aeon.base._base_series.VALID_INPUT_TYPES
+            A valid aeon time series data structure. See
+            aeon.base._base_series.VALID_INPUT_TYPES for aeon supported types.
+            Intended for algorithms which have another series as output.
+        axis: int or None
+            The time point axis of the input series if it is 2D. If ``axis==0``, it is
+            assumed each column is a time series and each row is a time point. i.e. the
+            shape of the data is ``(n_timepoints, n_channels)``. ``axis==1`` indicates
+            the time series are in rows, i.e. the shape of the data is
+            ``(n_channels, n_timepoints)``.
+            If None, the default class axis is used.
 
         Returns
         -------
-        params : dict or list of dict, default = {}
-            Parameters to create testing instances of the class.
+        Xt: one of aeon.base._base_series.VALID_INPUT_TYPES
+            New time series input reshaped to match the original input.
         """
-        # default parameters = empty dict
-        return {"axis": 0}
+        if axis is None:
+            axis = self.axis
+
+        # If a univariate only transformer, return a univariate series
+        if not self.get_tag("capability:multivariate"):
+            Xt = Xt.squeeze()
+
+        # return with input axis
+        if Xt.ndim == 1 or axis == self.axis:
+            return Xt
+        else:
+            return Xt.T
diff --git a/aeon/classification/compose/_ensemble.py b/aeon/classification/compose/_ensemble.py
index fbd5806f..eac060b8 100644
--- a/aeon/classification/compose/_ensemble.py
+++ b/aeon/classification/compose/_ensemble.py
@@ -32,9 +32,8 @@ class WeightedEnsembleClassifier(_HeterogenousMetaEstimator, BaseClassifier):
 
     Parameters
     ----------
-    classifiers : dict or None, default=None
-        Parameters for the ShapeletTransformClassifier module. If None, uses the
-        default parameters with a 2 hour transform contract.
+    classifiers : list of tuples (str, classifier) of aeon classifiers
+        Classifiers to apply to the input series.
     weights : float, or iterable of float, optional, default=None
         if float, ensemble weight for classifier i will be train score to this power
         if iterable of float, must be equal length as classifiers
diff --git a/aeon/clustering/_k_means.py b/aeon/clustering/_k_means.py
index 4b9d4a16..4aa04bd8 100644
--- a/aeon/clustering/_k_means.py
+++ b/aeon/clustering/_k_means.py
@@ -8,8 +8,8 @@ import numpy as np
 from numpy.random import RandomState
 from sklearn.utils import check_random_state
 
-from aeon.clustering.averaging import _resolve_average_callable
-from aeon.clustering.averaging._barycenter_averaging import VALID_BA_METRICS
+from aeon.clustering.averaging import VALID_BA_METRICS
+from aeon.clustering.averaging._averaging import _resolve_average_callable
 from aeon.clustering.base import BaseClusterer
 from aeon.distances import pairwise_distance
 
diff --git a/aeon/clustering/averaging/__init__.py b/aeon/clustering/averaging/__init__.py
index 03df3840..a5f3d2f6 100644
--- a/aeon/clustering/averaging/__init__.py
+++ b/aeon/clustering/averaging/__init__.py
@@ -1,5 +1,14 @@
 """Time series averaging metrics."""
 
-__all__ = ["elastic_barycenter_average", "mean_average", "_resolve_average_callable"]
-from aeon.clustering.averaging._averaging import _resolve_average_callable, mean_average
+__all__ = [
+    "elastic_barycenter_average",
+    "mean_average",
+    "petitjean_barycenter_average",
+    "subgradient_barycenter_average",
+    "VALID_BA_METRICS",
+]
+from aeon.clustering.averaging._averaging import mean_average
+from aeon.clustering.averaging._ba_petitjean import petitjean_barycenter_average
+from aeon.clustering.averaging._ba_subgradient import subgradient_barycenter_average
+from aeon.clustering.averaging._ba_utils import VALID_BA_METRICS
 from aeon.clustering.averaging._barycenter_averaging import elastic_barycenter_average
diff --git a/aeon/clustering/averaging/_ba_petitjean.py b/aeon/clustering/averaging/_ba_petitjean.py
new file mode 100644
index 00000000..69e1d352
--- /dev/null
+++ b/aeon/clustering/averaging/_ba_petitjean.py
@@ -0,0 +1,147 @@
+from typing import Optional, Tuple, Union
+
+import numpy as np
+from numba import njit
+
+from aeon.clustering.averaging._ba_utils import (
+    _get_alignment_path,
+    _get_init_barycenter,
+)
+
+
+def petitjean_barycenter_average(
+    X: np.ndarray,
+    distance: str = "dtw",
+    max_iters: int = 30,
+    tol=1e-5,
+    init_barycenter: Union[np.ndarray, str] = "mean",
+    precomputed_medoids_pairwise_distance: Optional[np.ndarray] = None,
+    verbose: bool = False,
+    random_state: Optional[int] = None,
+    **kwargs,
+) -> np.ndarray:
+    """Compute the barycenter average of time series using a elastic distance.
+
+    This implements an adapted version of 'petitjean' (original) DBA algorithm [1]_.
+
+    Parameters
+    ----------
+    X: np.ndarray, of shape (n_cases, n_channels, n_timepoints) or
+            (n_cases, n_timepoints)
+        A collection of time series instances to take the average from.
+    distance: str, default='dtw'
+        String defining the distance to use for averaging. Distance to
+        compute similarity between time series. A list of valid strings for metrics
+        can be found in the documentation form
+        :func:`aeon.distances.get_distance_function`.
+    max_iters: int, default=30
+        Maximum number iterations for dba to update over.
+    tol : float (default: 1e-5)
+        Tolerance to use for early stopping: if the decrease in cost is lower
+        than this value, the Expectation-Maximization procedure stops.
+    init_barycenter: np.ndarray or, default=None
+        The initial barycenter to use for the minimisation. If a np.ndarray is provided
+        it must be of shape ``(n_channels, n_timepoints)``. If a str is provided it must
+        be one of the following: ['mean', 'medoids', 'random'].
+    precomputed_medoids_pairwise_distance: np.ndarray (of shape (len(X), len(X)),
+                default=None
+        Precomputed medoids pairwise.
+    verbose: bool, default=False
+        Boolean that controls the verbosity.
+    random_state: int or None, default=None
+        Random state to use for the barycenter averaging.
+    **kwargs
+        Keyword arguments to pass to the distance metric.
+
+    Returns
+    -------
+    np.ndarray of shape (n_channels, n_timepoints)
+        Time series that is the average of the collection of instances provided.
+
+    References
+    ----------
+    .. [1] F. Petitjean, A. Ketterlin & P. Gancarski. A global averaging method
+       for dynamic time warping, with applications to clustering. Pattern
+       Recognition, Elsevier, 2011, Vol. 44, Num. 3, pp. 678-693
+    """
+    if len(X) <= 1:
+        return X
+
+    if X.ndim == 3:
+        _X = X
+    elif X.ndim == 2:
+        _X = X.reshape((X.shape[0], 1, X.shape[1]))
+    else:
+        raise ValueError("X must be a 2D or 3D array")
+
+    barycenter = _get_init_barycenter(
+        _X,
+        init_barycenter,
+        distance,
+        precomputed_medoids_pairwise_distance,
+        random_state,
+        **kwargs,
+    )
+
+    cost_prev = np.inf
+    if distance == "wdtw" or distance == "wddtw":
+        if "g" not in kwargs:
+            kwargs["g"] = 0.05
+    for i in range(max_iters):
+        barycenter, cost = _ba_one_iter_petitjean(barycenter, _X, distance, **kwargs)
+        if abs(cost_prev - cost) < tol:
+            break
+        elif cost_prev < cost:
+            break
+        else:
+            cost_prev = cost
+
+        if verbose:
+            print(f"[DBA] epoch {i}, cost {cost}")  # noqa: T001, T201
+    return barycenter
+
+
+@njit(cache=True, fastmath=True)
+def _ba_one_iter_petitjean(
+    barycenter: np.ndarray,
+    X: np.ndarray,
+    distance: str = "dtw",
+    window: Union[float, None] = None,
+    g: float = 0.0,
+    epsilon: Union[float, None] = None,
+    nu: float = 0.001,
+    lmbda: float = 1.0,
+    independent: bool = True,
+    c: float = 1.0,
+    descriptor: str = "identity",
+    reach: int = 30,
+    warp_penalty: float = 1.0,
+) -> Tuple[np.ndarray, float]:
+    X_size, X_dims, X_timepoints = X.shape
+    sum = np.zeros(X_timepoints)
+    alignment = np.zeros((X_dims, X_timepoints))
+    cost = 0.0
+    for i in range(X_size):
+        curr_ts = X[i]
+        curr_alignment, curr_cost = _get_alignment_path(
+            barycenter,
+            curr_ts,
+            distance,
+            window,
+            g,
+            epsilon,
+            nu,
+            lmbda,
+            independent,
+            c,
+            descriptor,
+            reach,
+            warp_penalty,
+        )
+
+        for j, k in curr_alignment:
+            alignment[:, k] += curr_ts[:, j]
+            sum[k] += 1
+        cost += curr_cost
+
+    return alignment / sum, cost
diff --git a/aeon/clustering/averaging/_ba_subgradient.py b/aeon/clustering/averaging/_ba_subgradient.py
new file mode 100644
index 00000000..67089c4b
--- /dev/null
+++ b/aeon/clustering/averaging/_ba_subgradient.py
@@ -0,0 +1,193 @@
+from typing import Optional, Union
+
+import numpy as np
+from numba import njit
+from sklearn.utils import check_random_state
+
+from aeon.clustering.averaging._ba_utils import (
+    _get_alignment_path,
+    _get_init_barycenter,
+)
+
+
+def subgradient_barycenter_average(
+    X: np.ndarray,
+    distance: str = "dtw",
+    max_iters: int = 30,
+    tol=1e-5,
+    init_barycenter: Union[np.ndarray, str] = "mean",
+    initial_step_size: float = 0.05,
+    final_step_size: float = 0.005,
+    precomputed_medoids_pairwise_distance: Optional[np.ndarray] = None,
+    verbose: bool = False,
+    random_state: Optional[int] = None,
+    **kwargs,
+) -> np.ndarray:
+    """Compute the stochastic subgradient barycenter average of time series.
+
+    Stochastic subgradient is much faster than petitjean, however, it is not guaranteed
+    to find the optimal solution.
+
+    This implements a stochastic subgradient DBA algorithm. This changes how
+    the average is computed. Unlike traditional methods, it calculates a subgradient
+    based on each individual time series within the dataset. The barycenter is then
+    iteratively updated using these subgradient. See [2]_ for more details.
+
+    Parameters
+    ----------
+    X: np.ndarray, of shape (n_cases, n_channels, n_timepoints) or
+            (n_cases, n_timepoints)
+        A collection of time series instances to take the average from.
+    distance: str, default='dtw'
+        String defining the distance to use for averaging. Distance to
+        compute similarity between time series. A list of valid strings for metrics
+        can be found in the documentation form
+        :func:`aeon.distances.get_distance_function`.
+    max_iters: int, default=30
+        Maximum number iterations for dba to update over.
+    tol : float (default: 1e-5)
+        Tolerance to use for early stopping: if the decrease in cost is lower
+        than this value, the Expectation-Maximization procedure stops.
+    init_barycenter: np.ndarray or, default=None
+        The initial barycenter to use for the minimisation. If a np.ndarray is provided
+        it must be of shape ``(n_channels, n_timepoints)``. If a str is provided it must
+        be one of the following: ['mean', 'medoids', 'random'].
+    initial_step_size : float (default: 0.05)
+        Initial step size for the subgradient descent algorithm.
+        Default value is suggested by [2]_.
+    final_step_size : float (default: 0.005)
+        Final step size for the subgradient descent algorithm.
+        Default value is suggested by [2]_.
+    precomputed_medoids_pairwise_distance: np.ndarray (of shape (len(X), len(X)),
+                default=None
+        Precomputed medoids pairwise.
+    verbose: bool, default=False
+        Boolean that controls the verbosity.
+    random_state: int or None, default=None
+        Random state to use for the barycenter averaging.
+    **kwargs
+        Keyword arguments to pass to the distance metric.
+
+    Returns
+    -------
+    np.ndarray of shape (n_channels, n_timepoints)
+        Time series that is the average of the collection of instances provided.
+
+    References
+    ----------
+    .. [1] F. Petitjean, A. Ketterlin & P. Gancarski. A global averaging method
+       for dynamic time warping, with applications to clustering. Pattern
+       Recognition, Elsevier, 2011, Vol. 44, Num. 3, pp. 678-693
+    .. [2] D. Schultz and B. Jain. Nonsmooth Analysis and Subgradient Methods
+       for Averaging in Dynamic Time Warping Spaces.
+       Pattern Recognition, 74, 340-358.
+    """
+    if len(X) <= 1:
+        return X
+
+    if X.ndim == 3:
+        _X = X
+    elif X.ndim == 2:
+        _X = X.reshape((X.shape[0], 1, X.shape[1]))
+    else:
+        raise ValueError("X must be a 2D or 3D array")
+
+    barycenter = _get_init_barycenter(
+        _X,
+        init_barycenter,
+        distance,
+        precomputed_medoids_pairwise_distance,
+        random_state,
+        **kwargs,
+    )
+
+    random_state = check_random_state(random_state)
+
+    cost_prev = np.inf
+    if distance == "wdtw" or distance == "wddtw":
+        if "g" not in kwargs:
+            kwargs["g"] = 0.05
+
+    current_step_size = initial_step_size
+    X_size = _X.shape[0]
+    for i in range(max_iters):
+        shuffled_indices = random_state.permutation(X_size)
+        barycenter, cost, current_step_size = _ba_one_iter_subgradient(
+            barycenter,
+            _X,
+            shuffled_indices,
+            distance,
+            initial_step_size,
+            final_step_size,
+            current_step_size,
+            i,
+            **kwargs,
+        )
+        if abs(cost_prev - cost) < tol:
+            break
+        elif cost_prev < cost:
+            break
+        else:
+            cost_prev = cost
+
+        if verbose:
+            print(f"[DBA] epoch {i}, cost {cost}")  # noqa: T001, T201
+    return barycenter
+
+
+@njit(cache=True, fastmath=True)
+def _ba_one_iter_subgradient(
+    barycenter: np.ndarray,
+    X: np.ndarray,
+    shuffled_indices: np.ndarray,
+    distance: str = "dtw",
+    initial_step_size: float = 0.05,
+    final_step_size: float = 0.005,
+    current_step_size: float = 0.05,
+    iteration: int = 0,
+    window: Union[float, None] = None,
+    g: float = 0.0,
+    epsilon: Union[float, None] = None,
+    nu: float = 0.001,
+    lmbda: float = 1.0,
+    independent: bool = True,
+    c: float = 1.0,
+    descriptor: str = "identity",
+    reach: int = 30,
+    warp_penalty: float = 1.0,
+):
+
+    X_size, X_dims, X_timepoints = X.shape
+    cost = 0.0
+    # Only update current_step_size on the first iteration
+    step_size_reduction = 0.0
+    if iteration == 0:
+        step_size_reduction = (initial_step_size - final_step_size) / X_size
+
+    for i in shuffled_indices:
+        curr_ts = X[i]
+        curr_alignment, curr_cost = _get_alignment_path(
+            barycenter,
+            X[i],
+            distance,
+            window,
+            g,
+            epsilon,
+            nu,
+            lmbda,
+            independent,
+            c,
+            descriptor,
+            reach,
+            warp_penalty,
+        )
+
+        new_ba = np.zeros((X_dims, X_timepoints))
+        for j, k in curr_alignment:
+            new_ba[:, k] += barycenter[:, k] - curr_ts[:, j]
+
+        barycenter -= (2.0 * current_step_size) * new_ba
+
+        current_step_size -= step_size_reduction
+        cost = curr_cost
+    return barycenter, cost, current_step_size
diff --git a/aeon/clustering/averaging/_ba_utils.py b/aeon/clustering/averaging/_ba_utils.py
new file mode 100644
index 00000000..1f40849c
--- /dev/null
+++ b/aeon/clustering/averaging/_ba_utils.py
@@ -0,0 +1,150 @@
+__maintainer__ = []
+
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+from numba import njit
+from sklearn.utils import check_random_state
+
+from aeon.distances import (
+    adtw_alignment_path,
+    ddtw_alignment_path,
+    dtw_alignment_path,
+    edr_alignment_path,
+    erp_alignment_path,
+    msm_alignment_path,
+    pairwise_distance,
+    shape_dtw_alignment_path,
+    twe_alignment_path,
+    wddtw_alignment_path,
+    wdtw_alignment_path,
+)
+
+
+def _medoids(
+    X: np.ndarray,
+    precomputed_pairwise_distance: Union[np.ndarray, None] = None,
+    distance: str = "dtw",
+    **kwargs,
+):
+    if X.shape[0] < 1:
+        return X
+
+    if precomputed_pairwise_distance is None:
+        precomputed_pairwise_distance = pairwise_distance(X, metric=distance, **kwargs)
+
+    x_size = X.shape[0]
+    distance_matrix = np.zeros((x_size, x_size))
+    for j in range(x_size):
+        for k in range(x_size):
+            distance_matrix[j, k] = precomputed_pairwise_distance[j, k]
+    return X[np.argmin(sum(distance_matrix))]
+
+
+def _get_init_barycenter(
+    X: np.ndarray,
+    init_barycenter: Union[np.ndarray, str],
+    distance: str,
+    precomputed_medoids_pw: Optional[np.ndarray] = None,
+    random_state: Optional[int] = None,
+    **kwargs,
+) -> np.ndarray:
+    if isinstance(init_barycenter, str):
+        if init_barycenter not in ["mean", "medoids", "random"]:
+            raise ValueError(
+                "init_barycenter string is invalid. Please use one of the "
+                "following: 'mean', 'medoids', 'random'"
+            )
+        if init_barycenter == "mean":
+            return X.mean(axis=0)
+        elif init_barycenter == "medoids":
+            return _medoids(
+                X,
+                precomputed_pairwise_distance=precomputed_medoids_pw,
+                distance=distance,
+                **kwargs,
+            )
+        else:
+            rng = check_random_state(random_state)
+            return X[rng.choice(X.shape[0])]
+    else:
+        if not isinstance(init_barycenter, np.ndarray):
+            raise ValueError(
+                "init_barycenter parameter is invalid. It must either be "
+                "a str or a np.ndarray"
+            )
+        if init_barycenter.shape != (X.shape[1], X.shape[2]):
+            if init_barycenter.ndim == 1:
+                return _get_init_barycenter(
+                    X=X,
+                    init_barycenter=init_barycenter.reshape(
+                        1, init_barycenter.shape[0]
+                    ),
+                    distance=distance,
+                    precomputed_medoids_pw=precomputed_medoids_pw,
+                    random_state=random_state,
+                    **kwargs,
+                )
+            raise ValueError(
+                f"init_barycenter shape is invalid. Expected {(X.shape[1], X.shape[2])}"
+                f" but got {init_barycenter.shape}"
+            )
+
+        return init_barycenter
+
+
+VALID_BA_METRICS = [
+    "dtw",
+    "ddtw",
+    "wdtw",
+    "wddtw",
+    "erp",
+    "edr",
+    "twe",
+    "msm",
+    "shape_dtw",
+]
+
+
+@njit(cache=True, fastmath=True)
+def _get_alignment_path(
+    center: np.ndarray,
+    ts: np.ndarray,
+    distance: str = "dtw",
+    window: Union[float, None] = None,
+    g: float = 0.0,
+    epsilon: Union[float, None] = None,
+    nu: float = 0.001,
+    lmbda: float = 1.0,
+    independent: bool = True,
+    c: float = 1.0,
+    descriptor: str = "identity",
+    reach: int = 30,
+    warp_penalty: float = 1.0,
+) -> Tuple[List[Tuple[int, int]], float]:
+    if distance == "dtw":
+        return dtw_alignment_path(ts, center, window)
+    elif distance == "ddtw":
+        return ddtw_alignment_path(ts, center, window)
+    elif distance == "wdtw":
+        return wdtw_alignment_path(ts, center, window, g)
+    elif distance == "wddtw":
+        return wddtw_alignment_path(ts, center, window, g)
+    elif distance == "erp":
+        return erp_alignment_path(ts, center, window, g)
+    elif distance == "edr":
+        return edr_alignment_path(ts, center, window, epsilon)
+    elif distance == "twe":
+        return twe_alignment_path(ts, center, window, nu, lmbda)
+    elif distance == "msm":
+        return msm_alignment_path(ts, center, window, independent, c)
+    elif distance == "shape_dtw":
+        return shape_dtw_alignment_path(
+            ts, center, window=window, descriptor=descriptor, reach=reach
+        )
+    elif distance == "adtw":
+        return adtw_alignment_path(ts, center, window=window, warp_penalty=warp_penalty)
+    else:
+        # When numba version > 0.57 add more informative error with what metric
+        # was passed.
+        raise ValueError("Distance parameter invalid")
diff --git a/aeon/clustering/averaging/_barycenter_averaging.py b/aeon/clustering/averaging/_barycenter_averaging.py
index 71ab6c7b..42eb0066 100644
--- a/aeon/clustering/averaging/_barycenter_averaging.py
+++ b/aeon/clustering/averaging/_barycenter_averaging.py
@@ -1,81 +1,84 @@
 __maintainer__ = []
 
-from typing import Optional, Tuple, Union
+from typing import Optional, Union
 
 import numpy as np
-from numba import njit
 
-from aeon.distances import (
-    adtw_alignment_path,
-    ddtw_alignment_path,
-    dtw_alignment_path,
-    edr_alignment_path,
-    erp_alignment_path,
-    msm_alignment_path,
-    pairwise_distance,
-    shape_dtw_alignment_path,
-    squared_distance,
-    twe_alignment_path,
-    wddtw_alignment_path,
-    wdtw_alignment_path,
-)
-
-
-def _medoids(
-    X: np.ndarray,
-    precomputed_pairwise_distance: Union[np.ndarray, None] = None,
-    distance: str = "dtw",
-    **kwargs,
-):
-    if X.shape[0] < 1:
-        return X
-
-    if precomputed_pairwise_distance is None:
-        precomputed_pairwise_distance = pairwise_distance(X, metric=distance, **kwargs)
-
-    x_size = X.shape[0]
-    distance_matrix = np.zeros((x_size, x_size))
-    for j in range(x_size):
-        for k in range(x_size):
-            distance_matrix[j, k] = precomputed_pairwise_distance[j, k]
-    return X[np.argmin(sum(distance_matrix))]
+from aeon.clustering.averaging._ba_petitjean import petitjean_barycenter_average
+from aeon.clustering.averaging._ba_subgradient import subgradient_barycenter_average
 
 
 def elastic_barycenter_average(
     X: np.ndarray,
     distance: str = "dtw",
     max_iters: int = 30,
-    tol=1e-5,
+    tol: float = 1e-5,
+    init_barycenter: Union[np.ndarray, str] = "mean",
+    method: str = "petitjean",
+    initial_step_size: float = 0.05,
+    final_step_size: float = 0.005,
     precomputed_medoids_pairwise_distance: Optional[np.ndarray] = None,
     verbose: bool = False,
+    random_state: Optional[int] = None,
     **kwargs,
 ) -> np.ndarray:
     """Compute the barycenter average of time series using a elastic distance.
 
-    This implements an adapted version of 'petitjean' (original) DBA algorithm [1]_.
+    This is a utility function that computes the barycenter average of a collection of
+    time series instances. The barycenter algorithm used can be select using the method
+    parameter. The following methods are available:
+    - 'petitjean': This implements an adapted version of 'petitjean' (original) DBA
+    algorithm [1]_.
+    - 'subgradient': This implements a stochastic subgradient DBA
+    algorithm [2]_.
+
+    Petitjean is slower but guaranteed to find the optimal solution. Stochastic
+    subgradient is much faster but not guaranteed to find the optimal solution.
+
+    For large datasets it is recommended to use the 'subgradient' method. This will
+    estimate the barycenter much faster than the 'petitjean' method. However, the
+    'subgradient' method is not guaranteed to find the optimal solution. If
+    computational time is not an issue, it is recommended to use the
+    'petitjean' method.
 
     Parameters
     ----------
     X: np.ndarray, of shape (n_cases, n_channels, n_timepoints) or
             (n_cases, n_timepoints)
         A collection of time series instances to take the average from.
-    distance: str or Callable, default='dtw'
+    distance: str, default='dtw'
         String defining the distance to use for averaging. Distance to
         compute similarity between time series. A list of valid strings for metrics
         can be found in the documentation form
         :func:`aeon.distances.get_distance_function`.
-        If Callable provided must be of the form (x, y) -> (float, np.ndarray)
-        where the first element is the distance and the second is the alignment path.
     max_iters: int, default=30
         Maximum number iterations for dba to update over.
     tol : float (default: 1e-5)
         Tolerance to use for early stopping: if the decrease in cost is lower
         than this value, the Expectation-Maximization procedure stops.
+    init_barycenter: np.ndarray or, default=None
+        The initial barycenter to use for the minimisation. If a np.ndarray is provided
+        it must be of shape ``(n_channels, n_timepoints)``. If a str is provided it must
+        be one of the following: ['mean', 'medoids', 'random'].
+        - 'mean': Uses mean of the time series instances.
+        - 'medoids': Uses medoids of the time series instances.
+        = 'random': Uses a random time series instance.
+    initial_step_size : float (default: 0.05)
+        Initial step size for the subgradient descent algorithm.
+        Default value is suggested by [2]_.
+    final_step_size : float (default: 0.005)
+        Final step size for the subgradient descent algorithm.
+        Default value is suggested by [2]_.
+    method: str, default='petitjean'
+        The method to use for the barycenter averaging. Valid strings are:
+        ['petitjean', 'subgradient'].
     precomputed_medoids_pairwise_distance: np.ndarray (of shape (len(X), len(X)),
                 default=None
         Precomputed medoids pairwise.
     verbose: bool, default=False
         Boolean that controls the verbosity.
+    random_state: int or None, default=None
+        Random state to use for the barycenter averaging.
     **kwargs
         Keyword arguments to pass to the distance metric.
 
@@ -89,104 +92,38 @@ def elastic_barycenter_average(
     .. [1] F. Petitjean, A. Ketterlin & P. Gancarski. A global averaging method
        for dynamic time warping, with applications to clustering. Pattern
        Recognition, Elsevier, 2011, Vol. 44, Num. 3, pp. 678-693
+    .. [2] D. Schultz and B. Jain. Nonsmooth Analysis and Subgradient Methods
+       for Averaging in Dynamic Time Warping Spaces.
+       Pattern Recognition, 74, 340-358.
     """
-    if len(X) <= 1:
-        return X
-
-    # center = X.mean(axis=0)
-    center = _medoids(
-        X,
-        distance=distance,
-        precomputed_pairwise_distance=precomputed_medoids_pairwise_distance,
-        **kwargs,
-    )
-
-    cost_prev = np.inf
-    if distance == "wdtw" or distance == "wddtw":
-        if "g" not in kwargs:
-            kwargs["g"] = 0.05
-    for i in range(max_iters):
-        center, cost = _ba_update(center, X, distance, **kwargs)
-        if abs(cost_prev - cost) < tol:
-            break
-        elif cost_prev < cost:
-            break
-        else:
-            cost_prev = cost
-
-        if verbose:
-            print(f"[DBA aeon] epoch {i}, cost {cost}")  # noqa: T001, T201
-    return center
-
-
-VALID_BA_METRICS = [
-    "dtw",
-    "ddtw",
-    "wdtw",
-    "wddtw",
-    "erp",
-    "edr",
-    "twe",
-    "msm",
-    "shape_dtw",
-]
-
-
-@njit(cache=True, fastmath=True)
-def _ba_update(
-    center: np.ndarray,
-    X: np.ndarray,
-    distance: str = "dtw",
-    window: Union[float, None] = None,
-    g: float = 0.0,
-    epsilon: Union[float, None] = None,
-    nu: float = 0.001,
-    lmbda: float = 1.0,
-    independent: bool = True,
-    c: float = 1.0,
-    descriptor: str = "identity",
-    reach: int = 30,
-    warp_penalty: float = 1.0,
-) -> Tuple[np.ndarray, float]:
-    X_size, X_dims, X_timepoints = X.shape
-    sum = np.zeros(X_timepoints)
-    alignment = np.zeros((X_dims, X_timepoints))
-    cost = 0.0
-    for i in range(X_size):
-        curr_ts = X[i]
-        if distance == "dtw":
-            curr_alignment, _ = dtw_alignment_path(curr_ts, center, window)
-        elif distance == "ddtw":
-            curr_alignment, _ = ddtw_alignment_path(curr_ts, center, window)
-        elif distance == "wdtw":
-            curr_alignment, _ = wdtw_alignment_path(curr_ts, center, window, g)
-        elif distance == "wddtw":
-            curr_alignment, _ = wddtw_alignment_path(curr_ts, center, window, g)
-        elif distance == "erp":
-            curr_alignment, _ = erp_alignment_path(curr_ts, center, window, g)
-        elif distance == "edr":
-            curr_alignment, _ = edr_alignment_path(curr_ts, center, window, epsilon)
-        elif distance == "twe":
-            curr_alignment, _ = twe_alignment_path(curr_ts, center, window, nu, lmbda)
-        elif distance == "msm":
-            curr_alignment, _ = msm_alignment_path(
-                curr_ts, center, window, independent, c
-            )
-        elif distance == "shape_dtw":
-            curr_alignment, _ = shape_dtw_alignment_path(
-                curr_ts, center, window=window, descriptor=descriptor, reach=reach
-            )
-        elif distance == "adtw":
-            curr_alignment, _ = adtw_alignment_path(
-                curr_ts, center, window=window, warp_penalty=warp_penalty
-            )
-        else:
-            # When numba version > 0.57 add more informative error with what metric
-            # was passed.
-            raise ValueError("Metric parameter invalid")
-        for j, k in curr_alignment:
-            alignment[:, k] += curr_ts[:, j]
-            sum[k] += 1
-            cost += squared_distance(curr_ts[:, j], center[:, k])
-
-    return alignment / sum, cost / X_timepoints
+    if method == "petitjean":
+        return petitjean_barycenter_average(
+            X,
+            distance=distance,
+            max_iters=max_iters,
+            tol=tol,
+            init_barycenter=init_barycenter,
+            precomputed_medoids_pairwise_distance=precomputed_medoids_pairwise_distance,
+            verbose=verbose,
+            random_state=random_state,
+            **kwargs,
+        )
+    elif method == "subgradient":
+        return subgradient_barycenter_average(
+            X,
+            distance=distance,
+            max_iters=max_iters,
+            tol=tol,
+            init_barycenter=init_barycenter,
+            initial_step_size=initial_step_size,
+            final_step_size=final_step_size,
+            precomputed_medoids_pairwise_distance=precomputed_medoids_pairwise_distance,
+            verbose=verbose,
+            random_state=random_state,
+            **kwargs,
+        )
+    else:
+        raise ValueError(
+            f"Invalid method: {method}. Please use one of the following: "
+            f"['petitjean', 'subgradient']"
+        )
diff --git a/aeon/segmentation/_clasp.py b/aeon/segmentation/_clasp.py
index eb44680f..65696ddf 100644
--- a/aeon/segmentation/_clasp.py
+++ b/aeon/segmentation/_clasp.py
@@ -225,6 +225,7 @@ class ClaSPSegmenter(BaseSegmenter):
         list
             List of change points found in X.
         """
+        X = X.squeeze()
         if len(X) - self.period_length < 2 * self.exclusion_radius * len(X):
             warnings.warn(
                 "Period-Length is larger than size of the time series", stacklevel=1
diff --git a/aeon/segmentation/_hmm.py b/aeon/segmentation/_hmm.py
index a8ab1b11..5d772a5b 100644
--- a/aeon/segmentation/_hmm.py
+++ b/aeon/segmentation/_hmm.py
@@ -363,6 +363,7 @@ class HMMSegmenter(BaseSegmenter):
         annotated_x : array-like, shape = [num_observations]
             Array of predicted class labels, same size as input.
         """
+        X = X.squeeze()
         self.num_states = len(self.emission_funcs)
         self.states = list(range(self.num_states))
         self.num_obs = len(X)
diff --git a/aeon/segmentation/base.py b/aeon/segmentation/base.py
index 9c265474..3ac70dad 100644
--- a/aeon/segmentation/base.py
+++ b/aeon/segmentation/base.py
@@ -7,8 +7,10 @@ from abc import ABC, abstractmethod
 from typing import List, final
 
 import numpy as np
+import pandas as pd
 
 from aeon.base import BaseSeriesEstimator
+from aeon.base._base_series import VALID_INPUT_TYPES
 
 
 class BaseSegmenter(BaseSeriesEstimator, ABC):
@@ -150,7 +152,7 @@ class BaseSegmenter(BaseSeriesEstimator, ABC):
         self.reset()
         if axis is None:  # If none given, assume it is correct.
             axis = self.axis
-        X = self._preprocess_series(X, axis)
+        X = self._preprocess_series(X, axis, True)
         if y is not None:
             y = self._check_y(y)
         self._fit(X=X, y=y)
@@ -184,7 +186,7 @@ class BaseSegmenter(BaseSeriesEstimator, ABC):
         self.check_is_fitted()
         if axis is None:
             axis = self.axis
-        X = self._preprocess_series(X, axis)
+        X = self._preprocess_series(X, axis, self.get_class_tag("fit_is_empty"))
         return self._predict(X)
 
     def fit_predict(self, X, y=None, axis=None):
@@ -203,6 +205,48 @@ class BaseSegmenter(BaseSeriesEstimator, ABC):
         """Create and return a segmentation of X."""
         ...
 
+    def _check_y(self, y: VALID_INPUT_TYPES):
+        """Check y specific to segmentation.
+
+        y must be a univariate series
+        """
+        if type(y) not in VALID_INPUT_TYPES:
+            raise ValueError(
+                f"Error in input type for y: it should be one of "
+                f"{VALID_INPUT_TYPES}, saw {type(y)}"
+            )
+        if isinstance(y, np.ndarray):
+            # Check valid shape
+            if y.ndim > 1:
+                raise ValueError(
+                    "Error in input type for y: y input as np.ndarray " "should be 1D"
+                )
+            if not (
+                issubclass(y.dtype.type, np.integer)
+                or issubclass(y.dtype.type, np.floating)
+            ):
+                raise ValueError(
+                    "Error in input type for y: y input must contain " "floats or ints"
+                )
+        elif isinstance(y, pd.Series):
+            if not pd.api.types.is_numeric_dtype(y):
+                raise ValueError(
+                    "Error in input type for y: y input as pd.Series must be numeric"
+                )
+        else:  # pd.DataFrame
+            if y.shape[1] > 2:
+                raise ValueError(
+                    "Error in input type for y: y input as pd.DataFrame "
+                    "should have a single "
+                    "column series"
+                )
+
+            if not all(pd.api.types.is_numeric_dtype(y[col]) for col in y.columns):
+                raise ValueError(
+                    "Error in input type for y: y input as pd.DataFrame "
+                    "must be numeric"
+                )
+
     @classmethod
     def to_classification(cls, change_points: List[int], length: int):
         """Convert change point locations to a classification vector.
diff --git a/aeon/transformations/series/_matrix_profile.py b/aeon/transformations/series/_matrix_profile.py
index 965d8f77..9cc1d6ce 100644
--- a/aeon/transformations/series/_matrix_profile.py
+++ b/aeon/transformations/series/_matrix_profile.py
@@ -67,6 +67,8 @@ class MatrixProfileSeriesTransformer(BaseSeriesTransformer):
         _check_soft_dependencies("stumpy", severity="error")
         import stumpy
 
+        X = X.squeeze()
+
         self.matrix_profile_ = stumpy.stump(X, self.window_length)
         return self.matrix_profile_[:, 0].astype("float")
 
diff --git a/aeon/transformations/series/base.py b/aeon/transformations/series/base.py
index 878dc2c6..209a2984 100644
--- a/aeon/transformations/series/base.py
+++ b/aeon/transformations/series/base.py
@@ -70,7 +70,7 @@ class BaseSeriesTransformer(BaseSeriesEstimator, BaseTransformer, metaclass=ABCM
         self.reset()
         if axis is None:  # If none given, assume it is correct.
             axis = self.axis
-        X = self._preprocess_series(X, axis=axis)
+        X = self._preprocess_series(X, axis=axis, store_metadata=True)
         if y is not None:
             self._check_y(y)
         self._fit(X=X, y=y)
@@ -109,7 +109,9 @@ class BaseSeriesTransformer(BaseSeriesEstimator, BaseTransformer, metaclass=ABCM
 
         if axis is None:
             axis = self.axis
-        X = self._preprocess_series(X, axis=axis)
+        X = self._preprocess_series(
+            X, axis=axis, store_metadata=self.get_class_tag("fit_is_empty")
+        )
         Xt = self._transform(X)
         return self._postprocess_series(Xt, axis=axis)
 
@@ -145,7 +147,7 @@ class BaseSeriesTransformer(BaseSeriesEstimator, BaseTransformer, metaclass=ABCM
         """
         # input checks and datatype conversion, to avoid doing in both fit and transform
         self.reset()
-        X = self._preprocess_series(X, axis=axis)
+        X = self._preprocess_series(X, axis=axis, store_metadata=True)
         Xt = self._fit_transform(X=X, y=y)
         self._is_fitted = True
         return self._postprocess_series(Xt, axis=axis)
@@ -179,7 +181,9 @@ class BaseSeriesTransformer(BaseSeriesEstimator, BaseTransformer, metaclass=ABCM
 
         # check whether is fitted
         self.check_is_fitted()
-        X = self._preprocess_series(X, axis=axis)
+        X = self._preprocess_series(
+            X, axis=axis, store_metadata=self.get_class_tag("fit_is_empty")
+        )
         Xt = self._inverse_transform(X=X, y=y)
         return self._postprocess_series(Xt, axis=axis)
 
@@ -204,7 +208,7 @@ class BaseSeriesTransformer(BaseSeriesEstimator, BaseTransformer, metaclass=ABCM
         """
         # check whether is fitted
         self.check_is_fitted()
-        X = self._preprocess_series(X, axis=axis)
+        X = self._preprocess_series(X, axis, self.get_class_tag("fit_is_empty"))
         return self._update(X=X, y=y, update_params=update_params)
 
     def _fit(self, X, y=None):
diff --git a/examples/base/series_estimator.ipynb b/examples/base/series_estimator.ipynb
index 91e63d4b..50702eca 100644
--- a/examples/base/series_estimator.ipynb
+++ b/examples/base/series_estimator.ipynb
@@ -5,198 +5,458 @@
    "source": [
     "# BaseSeriesEstimator\n",
     "\n",
-    "The ``BaseSeriesEstimator``class is a base class for estimators that take a single\n",
-    "series as input rather than a collection of time series (see\n",
-    "``BaseCollectionEstimator``). This notebook describes the major design issues to bare\n",
-    " in mind if using any class that inherits from ``BaseSeriesEstimator``. To use any\n",
-    " base estimator all you need to understand is the meaning of ``axis`` and the capability\n",
-    "  tags.\n",
+    "The `BaseSeriesEstimator` class is a base class for estimators that take a single series (both univariate and multivariate) as input rather than a collection of time series (see `BaseCollectionEstimator`). This notebook describes the major design issues to keep in mind if using any class that inherits from `BaseSeriesEstimator`. These are:\n",
+    "- `BaseSeriesTransformer` for single series transformations\n",
+    "- `BaseSegmenter` for segmentation\n",
+    "- `BaseAnomalyDetector` for anomaly detection\n",
     "\n",
-    "``BaseSeriesEstimator`` handles the preprocessing required for a single before being\n",
-    "used in a method such as ``fit``. These are used in base classes by\n",
-    "applying the protected method ``preprocess_series``. The key steps to note are:\n",
-    "1. Input data type should be a ``np.ndarray``, a ``pd.Series`` or a ``pd.DataFrame``.\n",
-    "2. The input data will be transformed into the type required by the estimator as\n",
-    "determined by the tag ``X_inner_type``.\n",
-    "3. If the estimator can only work with univariate time series\n",
-    "(``capability:multivariate`` set to False) then the input data will be converted to a\n",
-    " 1D numpy array or a pandas Series.\n",
-    "4. If the estimator has the ability to handle multivariate time series as determined\n",
-    " by the tag ``capability:multivariate``, then the input data will stored in either a 2\n",
-    "  D numpy array or a pandas DataFrames.\n",
-    "5. If the data is multivariate, then the ``axis`` varaible of the estimator controls\n",
-    "how it is interpreted. If ``axis==0`` then each column is a time series, and\n",
-    "each row is a time point: i.e. the shape of the data is ``(n_timepoints,n_channels)``\n",
-    ". If ``axis == 1`` indicates the time series are in rows, i.e. the shape of the data\n",
-    "is ``(n_channels, n_timepoints)``.\n",
+    "To use any algorithms extending the base estimator all you need to understand is the meaning of the `axis` parameter and the capability tags. `BaseSeriesEstimator` handles the preprocessing required before being used in methods such as `fit` and `predict`. These are used in inheriting base classes by applying the protected method `_preprocess_series`. The key steps to note are:\n",
+    "1. Input data type should be a `np.ndarray`, a `pd.Series` or a `pd.DataFrame`.\n",
+    "2. Unless the `X_inner_type` of the estimator is `pd.Series`, the `axis` variable of the estimator controls how the input data is interpreted in methods such as `fit`, `predict` and `transform`. If `axis==0` then each column is a time series, and each row is a time point: i.e. the shape of the input data is `(n_timepoints, n_channels)`. If `axis==1` indicates the time series are in rows, i.e. the shape of the data is `(n_channels, n_timepoints)`. It is important to set this correctly or check the default used, otherwise your data may be processed incorrectly.\n",
+    "3. The input data will be transformed into the type required by the estimator as determined by the tag `X_inner_type`. This will also reshape the array to use the correct time point `axis` and expand the input to 2D if it is a 1D `np.ndarray`.\n",
+    "4. If the estimator can only work with univariate time series (`capability:multivariate` set to False) then the input data shape must be 1D or 2D with the selected channel `axis` being size 1.\n",
+    "5. If the estimator can only work with multivariate time series (`capability:univariate` set to False), then the input data must be 2D, with the selected channel `axis` greater than 1. `pd.Series` is not supported in this case.\n",
     "\n",
-    "We demonstrate this with calls to private methods. This is purely to aide understanding\n",
-    " and should not be used in practice."
+    "We demonstrate this with calls to private methods. This is purely to aid understanding and should not be used in practice."
    ],
    "metadata": {
     "collapsed": false
    }
   },
   {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-10T15:04:57.444241Z",
+     "start_time": "2024-04-10T15:04:55.744636Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "\n",
+    "from aeon.base import BaseSeriesEstimator\n",
+    "\n",
+    "# We use the abstract base class for example purposes, regular classes will not\n",
+    "# have a class axis parameter.\n",
+    "bs = BaseSeriesEstimator(axis=0)"
+   ],
+   "outputs": [],
+   "execution_count": 1
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": "## Univariate examples"
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-10T15:04:57.453217Z",
+     "start_time": "2024-04-10T15:04:57.445239Z"
+    }
+   },
    "cell_type": "code",
-   "execution_count": 28,
+   "source": [
+    "# By default, \"capability:multivariate\" is False, axis is 0 and X_inner_type is\n",
+    "# np.ndarray\n",
+    "# With this config, the output should always be an np.ndarray shape (100, 1)\n",
+    "d1 = np.random.random(size=(100))\n",
+    "d2 = bs._preprocess_series(d1, axis=0, store_metadata=True)\n",
+    "print(\n",
+    "    f\"1. Input type = {type(d1)}, input shape = {d1.shape}, \"\n",
+    "    f\"output type = {type(d2)}, output shape = {d2.shape}\"\n",
+    ")"
+   ],
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "1. Input shape =  (100,)  output type =  <class 'numpy.ndarray'>  output shape =  (100,)\n",
-      "2. Input shape =  (1, 100)  output type =  <class 'numpy.ndarray'>  output shape =  (100,)\n",
-      "3. Input shape =  (100,)  output type =  <class 'numpy.ndarray'>  output shape =  (100,)\n",
-      "4. Input shape =  (100,)  output type =  <class 'numpy.ndarray'>  output shape =  (100,)\n",
-      "5. Input shape =  (100, 1)  output type =  <class 'numpy.ndarray'>  output shape =  (100,)\n"
+      "1. Input type = <class 'numpy.ndarray'>, input shape = (100,), output type = <class 'numpy.ndarray'>, output shape = (100, 1)\n"
      ]
     }
    ],
+   "execution_count": 2
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-10T15:04:57.461197Z",
+     "start_time": "2024-04-10T15:04:57.455212Z"
+    }
+   },
+   "cell_type": "code",
    "source": [
-    "# Univariate examples\n",
-    "import numpy as np\n",
-    "import pandas as pd\n",
-    "import pytest\n",
-    "\n",
-    "from aeon.base import BaseSeriesEstimator\n",
-    "\n",
-    "bs = BaseSeriesEstimator()\n",
-    "# By default, \"capability:multivariate\" is False, axis is 0 and\n",
-    "# X_inner_type is np.ndarray\n",
+    "# The axis parameter will not change the output shape of 1D inputs such as pd.Series\n",
+    "# or univariate np.ndarray\n",
     "d1 = np.random.random(size=(100))\n",
-    "# With this config, the output should always be an np.ndarray\n",
-    "# shape (100,)\n",
-    "d2 = bs._preprocess_series(d1, axis=0)\n",
+    "d2 = bs._preprocess_series(d1, axis=1, store_metadata=True)\n",
     "print(\n",
-    "    \"1. Input shape = \",\n",
-    "    d1.shape,\n",
-    "    \" output type = \",\n",
-    "    type(d2),\n",
-    "    \" output shape = \",\n",
-    "    d2.shape,\n",
-    ")\n",
-    "# 2D numpy shape (m,1) or (1,m) get converted to 1D numpy array\n",
-    "# if multivariate is False\n",
+    "    f\"2. Input type = {type(d1)}, input shape = {d1.shape}, \"\n",
+    "    f\"output type = {type(d2)}, output shape = {d2.shape}\"\n",
+    ")"
+   ],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2. Input type = <class 'numpy.ndarray'>, input shape = (100,), output type = <class 'numpy.ndarray'>, output shape = (100, 1)\n"
+     ]
+    }
+   ],
+   "execution_count": 3
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-10T15:04:57.468177Z",
+     "start_time": "2024-04-10T15:04:57.462193Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "# A 2D array with the channel axis of size 1 will produce the same result\n",
+    "d1 = np.random.random(size=(100, 1))\n",
+    "d2 = bs._preprocess_series(d1, axis=0, store_metadata=True)\n",
+    "print(\n",
+    "    f\"3. Input type = {type(d1)}, input shape = {d1.shape}, \"\n",
+    "    f\"output type = {type(d2)}, output shape = {d2.shape}\"\n",
+    ")"
+   ],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "3. Input type = <class 'numpy.ndarray'>, input shape = (100, 1), output type = <class 'numpy.ndarray'>, output shape = (100, 1)\n"
+     ]
+    }
+   ],
+   "execution_count": 4
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-10T15:04:57.474979Z",
+     "start_time": "2024-04-10T15:04:57.470430Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "# The shape used can be swapped, but the axis parameter must be set correctly\n",
     "d1 = np.random.random(size=(1, 100))\n",
-    "d2 = bs._preprocess_series(d1, axis=0)\n",
+    "d2 = bs._preprocess_series(d1, axis=1, store_metadata=True)\n",
     "print(\n",
-    "    \"2. Input shape = \",\n",
-    "    d1.shape,\n",
-    "    \" output type = \",\n",
-    "    type(d2),\n",
-    "    \" output shape = \",\n",
-    "    d2.shape,\n",
-    ")\n",
+    "    f\"4. Input type = {type(d1)}, input shape = {d1.shape}, \"\n",
+    "    f\"output type = {type(d2)}, output shape = {d2.shape}\"\n",
+    ")"
+   ],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "4. Input type = <class 'numpy.ndarray'>, input shape = (1, 100), output type = <class 'numpy.ndarray'>, output shape = (100, 1)\n"
+     ]
+    }
+   ],
+   "execution_count": 5
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-10T15:04:57.481960Z",
+     "start_time": "2024-04-10T15:04:57.475976Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "# Other types will be converted to X_inner_type\n",
     "d1 = pd.Series(np.random.random(size=(100)))\n",
-    "d2 = bs._preprocess_series(d1, axis=0)\n",
+    "d2 = bs._preprocess_series(d1, axis=0, store_metadata=True)\n",
     "print(\n",
-    "    \"3. Input shape = \",\n",
-    "    d1.shape,\n",
-    "    \" output type = \",\n",
-    "    type(d2),\n",
-    "    \" output shape = \",\n",
-    "    d2.shape,\n",
-    ")\n",
-    "# Axis is irrelevant for univariate data\n",
-    "d2 = bs._preprocess_series(d1, axis=1)\n",
-    "print(\n",
-    "    \"4. Input shape = \",\n",
-    "    d1.shape,\n",
-    "    \" output type = \",\n",
-    "    type(d2),\n",
-    "    \" output shape = \",\n",
-    "    d2.shape,\n",
-    ")\n",
+    "    f\"5. Input type = {type(d1)}, input shape = {d1.shape}, \"\n",
+    "    f\"output type = {type(d2)}, output shape = {d2.shape}\"\n",
+    ")"
+   ],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "5. Input type = <class 'pandas.core.series.Series'>, input shape = (100,), output type = <class 'numpy.ndarray'>, output shape = (100, 1)\n"
+     ]
+    }
+   ],
+   "execution_count": 6
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-10T15:04:57.489939Z",
+     "start_time": "2024-04-10T15:04:57.482958Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
     "d1 = pd.DataFrame(np.random.random(size=(100, 1)))\n",
-    "d2 = bs._preprocess_series(d1, axis=0)\n",
+    "d2 = bs._preprocess_series(d1, axis=0, store_metadata=True)\n",
     "print(\n",
-    "    \"5. Input shape = \",\n",
-    "    d1.shape,\n",
-    "    \" output type = \",\n",
-    "    type(d2),\n",
-    "    \" output shape = \",\n",
-    "    d2.shape,\n",
-    ")\n",
-    "\n",
-    "# Passing a multivariate array will raise an error\n",
-    "with pytest.raises(ValueError, match=r\"Multivariate data not supported\"):\n",
-    "    bs._check_X(np.random.random(size=(4, 100)))"
+    "    f\"6. Input type = {type(d1)}, input shape = {d1.shape}, \"\n",
+    "    f\"output type = {type(d2)}, output shape = {d2.shape}\"\n",
+    ")"
    ],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "6. Input type = <class 'pandas.core.frame.DataFrame'>, input shape = (100, 1), output type = <class 'numpy.ndarray'>, output shape = (100, 1)\n"
+     ]
+    }
+   ],
+   "execution_count": 7
+  },
+  {
    "metadata": {
-    "collapsed": false
-   }
+    "ExecuteTime": {
+     "end_time": "2024-04-10T15:04:57.497917Z",
+     "start_time": "2024-04-10T15:04:57.490936Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "bs = bs.set_tags(**{\"X_inner_type\": \"pd.Series\"})\n",
+    "d1 = np.random.random(size=(100))\n",
+    "d2 = bs._preprocess_series(d1, axis=0, store_metadata=True)\n",
+    "print(\n",
+    "    f\"7. Input type = {type(d1)}, input shape = {d1.shape}, \"\n",
+    "    f\"output type = {type(d2)}, output shape = {d2.shape}\"\n",
+    ")\n",
+    "bs = bs.set_tags(**{\"X_inner_type\": \"np.ndarray\"})"
+   ],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "7. Input type = <class 'numpy.ndarray'>, input shape = (100,), output type = <class 'pandas.core.series.Series'>, output shape = (100,)\n"
+     ]
+    }
+   ],
+   "execution_count": 8
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "source": [
+    "# Passing a multivariate array will raise an error if capability:multivariate is False\n",
+    "d1 = np.random.random(size=(100, 5))\n",
+    "try:\n",
+    "    bs._preprocess_series(d1, axis=0, store_metadata=True)\n",
+    "except ValueError as e:\n",
+    "    print(f\"8. {e}\")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-04-10T15:04:57.504899Z",
+     "start_time": "2024-04-10T15:04:57.498915Z"
+    }
+   },
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "1. Input shape =  (4, 100)  output type =  <class 'numpy.ndarray'>  output shape =  (4, 100)\n",
-      "2. Input shape =  (4, 100)  output type =  <class 'numpy.ndarray'>  output shape =  (100, 4)\n",
-      "2. Input type = <class 'pandas.core.frame.DataFrame'> Input shape =  (4, 100)  output type =  <class 'numpy.ndarray'> output shape =  (100, 4)\n"
+      "8. Multivariate data not supported by BaseSeriesEstimator\n"
      ]
     }
    ],
+   "execution_count": 9
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": "## Multivariate examples"
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-10T15:04:57.510883Z",
+     "start_time": "2024-04-10T15:04:57.505896Z"
+    }
+   },
+   "cell_type": "code",
    "source": [
-    "# Multivariate examples\n",
-    "# Set tags\n",
-    "bs.set_tags(**{\"capability:multivariate\": True})\n",
-    "d1 = np.random.random(size=(4, 100))\n",
-    "# Axis 0 means each row is a time series\n",
-    "d2 = bs._preprocess_series(d1, axis=0)\n",
+    "# The capability:multivariate tag must be set to True to work with multivariate series\n",
+    "# If the estimator does not have this tag, then the implementation cannot handle the\n",
+    "# input\n",
+    "bs = bs.set_tags(**{\"capability:multivariate\": True})\n",
+    "# Both of these can be True at the same time, but for examples sake we disable\n",
+    "# univariate\n",
+    "bs = bs.set_tags(**{\"capability:univariate\": False})"
+   ],
+   "outputs": [],
+   "execution_count": 10
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "# axis 0 means each row is a time series\n",
+    "d1 = np.random.random(size=(100, 5))\n",
+    "d2 = bs._preprocess_series(d1, axis=0, store_metadata=True)\n",
     "print(\n",
-    "    \"1. Input shape = \",\n",
-    "    d1.shape,\n",
-    "    \" output type = \",\n",
-    "    type(d2),\n",
-    "    \" output shape = \",\n",
-    "    d2.shape,\n",
+    "    f\"1. Input type = {type(d1)}, input shape = {d1.shape}, \"\n",
+    "    f\"output type = {type(d2)}, output shape = {d2.shape}\"\n",
     ")\n",
-    "# Axis 1 means each column is a time series\n",
-    "d2 = bs._preprocess_series(d1, axis=1)\n",
+    "n_channels = bs.metadata_[\"n_channels\"]\n",
+    "print(f\"n_channels: {n_channels}\")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-04-10T15:04:57.517865Z",
+     "start_time": "2024-04-10T15:04:57.511880Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1. Input type = <class 'numpy.ndarray'>, input shape = (100, 5), output type = <class 'numpy.ndarray'>, output shape = (100, 5)\n",
+      "n_channels: 5\n"
+     ]
+    }
+   ],
+   "execution_count": 11
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-10T15:04:57.524846Z",
+     "start_time": "2024-04-10T15:04:57.518861Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "# axis 1 means each column is a time series. If the axis is set incorrectly, the\n",
+    "# output shape will be wrong\n",
+    "d1 = np.random.random(size=(100, 5))\n",
+    "d2 = bs._preprocess_series(d1, axis=1, store_metadata=True)\n",
     "print(\n",
-    "    \"2. Input shape = \",\n",
-    "    d1.shape,\n",
-    "    \" output type = \",\n",
-    "    type(d2),\n",
-    "    \" output shape = \",\n",
-    "    d2.shape,\n",
+    "    f\"2. Input type = {type(d1)}, input shape = {d1.shape}, \"\n",
+    "    f\"output type = {type(d2)}, output shape = {d2.shape}\"\n",
     ")\n",
-    "d1 = pd.DataFrame(d1)\n",
-    "d2 = bs._preprocess_series(d1, axis=1)\n",
+    "n_channels = bs.metadata_[\"n_channels\"]\n",
+    "print(f\"n_channels: {n_channels}\")"
+   ],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2. Input type = <class 'numpy.ndarray'>, input shape = (100, 5), output type = <class 'numpy.ndarray'>, output shape = (5, 100)\n",
+      "n_channels: 100\n"
+     ]
+    }
+   ],
+   "execution_count": 12
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-10T15:04:57.532825Z",
+     "start_time": "2024-04-10T15:04:57.525844Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "# Conversions work similar to univariate series, but there is more emphasis on correctly\n",
+    "# setting the axis parameter\n",
+    "d1 = pd.DataFrame(np.random.random(size=(100, 5)))\n",
+    "d2 = bs._preprocess_series(d1, axis=0, store_metadata=True)\n",
     "print(\n",
-    "    \"2. Input type =\",\n",
-    "    type(d1),\n",
-    "    \"Input shape = \",\n",
-    "    d1.shape,\n",
-    "    \" output type = \",\n",
-    "    type(d2),\n",
-    "    \"output shape = \",\n",
-    "    d2.shape,\n",
-    ")"
+    "    f\"3. Input type = {type(d1)}, input shape = {d1.shape}, \"\n",
+    "    f\"output type = {type(d2)}, output shape = {d2.shape}\"\n",
+    ")\n",
+    "n_channels = bs.metadata_[\"n_channels\"]\n",
+    "print(f\"n_channels: {n_channels}\")"
    ],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "3. Input type = <class 'pandas.core.frame.DataFrame'>, input shape = (100, 5), output type = <class 'numpy.ndarray'>, output shape = (100, 5)\n",
+      "n_channels: 5\n"
+     ]
+    }
+   ],
+   "execution_count": 13
+  },
+  {
    "metadata": {
-    "collapsed": false
-   }
+    "ExecuteTime": {
+     "end_time": "2024-04-10T15:04:57.540803Z",
+     "start_time": "2024-04-10T15:04:57.535816Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "# As pd.Series is univariate only, it is not allowed as an inner type for multivariate\n",
+    "# This should not be an issue for usage, just for development\n",
+    "bs = bs.set_tags(**{\"X_inner_type\": \"pd.Series\"})\n",
+    "d1 = np.random.random(size=(100, 5))\n",
+    "try:\n",
+    "    d2 = bs._preprocess_series(d1, axis=1, store_metadata=True)\n",
+    "except ValueError as e:\n",
+    "    print(f\"4. {e}\")\n",
+    "bs = bs.set_tags(**{\"X_inner_type\": \"np.ndarray\"})"
+   ],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "4. Cannot convert to pd.Series for multivariate capable estimators\n"
+     ]
+    }
+   ],
+   "execution_count": 14
   },
   {
-   "cell_type": "markdown",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-10T15:04:57.547785Z",
+     "start_time": "2024-04-10T15:04:57.541801Z"
+    }
+   },
+   "cell_type": "code",
    "source": [
-    "If implementing a new estimator that extends ``BaseSeriesEstimator`` then just set\n",
-    "the ``axis`` to the shape you want to work with by passing it to the\n",
-    "``BaseSeriesEstimator`` constructor. If your estimator can handle\n",
-    "multivariate series, set the tag and set the ``capability:multivariate`` tag to True.\n",
-    " The data will always then be passed to your estimator in ``(n_channels,\n",
-    " n_timepoints)`` if axis is 1, or ``(n_timepoints, n_channels)``\n",
-    " if axis is 0, either\n",
-    "  in numpy arrays or pandas DataFrame, dependning on ``X_inner_type`` tag. If\n",
-    " a univariate series is passed  it will be passed in ``(1, n_timepoints)`` if axis is\n",
-    "  0, or ``(n_timepoints, 1)`` if the estimator axis is 0."
+    "# Passing a univariate array will raise an error if capability:univariate is False\n",
+    "d1 = pd.Series(np.random.random(size=(100,)))\n",
+    "try:\n",
+    "    d2 = bs._preprocess_series(d1, axis=0, store_metadata=True)\n",
+    "except ValueError as e:\n",
+    "    print(f\"5. {e}\")"
+   ],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "5. Univariate data not supported by BaseSeriesEstimator\n"
+     ]
+    }
    ],
+   "execution_count": 15
+  },
+  {
+   "cell_type": "markdown",
+   "source": "If implementing a new estimator that extends `BaseSeriesEstimator` then just set the `axis` to the shape you want to work with by passing it to the `BaseSeriesEstimator` constructor. If your estimator can handle multivariate series, set the tag and set the `capability:multivariate` tag to `True`.  Set the `X_inner_type` tag if you wish to use a datatype other than `np.ndarray`.",
    "metadata": {
     "collapsed": false
    }


Test patch:
diff --git a/aeon/base/tests/test_base_series.py b/aeon/base/tests/test_base_series.py
index bbe816f8..e0ea7a40 100644
--- a/aeon/base/tests/test_base_series.py
+++ b/aeon/base/tests/test_base_series.py
@@ -3,81 +3,46 @@
 import numpy as np
 import pandas as pd
 import pytest
+from numpy.testing import assert_equal
 
 from aeon.base import BaseSeriesEstimator
 
-
-def test__check_X():
-    """Test if capabilities correctly tested."""
-    dummy1 = BaseSeriesEstimator()
-    invalid_always = np.array(["1", "2", "3", "4", "5"])
-    with pytest.raises(ValueError, match=r"Error in input type should be one of"):
-        dummy1._check_X("String input")
-    with pytest.raises(ValueError, match=r"array must contain floats or ints"):
-        dummy1._check_X(invalid_always)
-    with pytest.raises(ValueError, match=r"pd.Series must be numeric"):
-        dummy1._check_X(pd.Series(invalid_always))
-    with pytest.raises(ValueError, match=r"pd.DataFrame must be numeric"):
-        dummy1._check_X(pd.DataFrame(invalid_always))
-    multi_np = np.random.random(size=(5, 10))
-    multi_pd = pd.DataFrame(multi_np)
-    with pytest.raises(ValueError, match=r"Multivariate data not supported"):
-        dummy1._check_X(multi_np)
-    with pytest.raises(ValueError, match=r"Multivariate data not supported"):
-        dummy1._check_X(multi_pd)
-    missing = np.random.rand(4, 10)
-    missing[3][8] = np.NAN
-    with pytest.raises(ValueError, match=r"Missing values not supported"):
-        dummy1._check_X(missing)
-
-    dummy2 = BaseSeriesEstimator()
-    all_tags = {
-        "capability:multivariate": True,
-        "capability:missing_values": True,
-    }
-    dummy2.set_tags(**all_tags)
-    meta = dummy2._check_X(multi_np)
-    meta2 = dummy2._check_X(multi_pd)
-    assert meta == meta2
-    assert meta["multivariate"]
-    assert not meta["missing_values"]
-    multi_np[0][0] = np.NAN
-    multi_pd[0][0] = np.NAN
-    uni_missing = pd.Series(np.array([1.0, np.NAN, 2.0]))
-    meta = dummy2._check_X(multi_np)
-    meta2 = dummy2._check_X(multi_pd)
-    meta3 = dummy2._check_X(uni_missing)
-    assert (
-        meta["missing_values"] and meta2["missing_values"] and meta3["missing_values"]
-    )
-    work_always = np.random.random(size=(10))
-    meta = dummy1._check_X(work_always)
-    meta2 = dummy1._check_X(pd.Series(work_always))
-    meta3 = dummy1._check_X(pd.DataFrame(work_always))
-    assert meta == meta2 == meta3
-    all_tags = {"capability:univariate": False}
-    dummy2.set_tags(**all_tags)
-    uni_np = np.random.random(size=(10))
-    uni_pd = pd.Series(uni_np)
-    with pytest.raises(ValueError, match=r"Univariate data not supported"):
-        dummy2._check_X(uni_np)
-    with pytest.raises(ValueError, match=r"Univariate data not supported"):
-        dummy2._check_X(uni_pd)
-    wrong_size = np.random.rand(10, 2, 20)
-    with pytest.raises(ValueError, match=r"Should be 1D or 2D"):
-        dummy2._check_X(wrong_size)
-
-
 UNIVARIATE = {
     "np.ndarray": np.random.random(size=(20)),
     "pd.Series": pd.Series(np.random.random(size=(20))),
-    "pd.DataFrame": pd.DataFrame(np.random.random(size=(20))),
+    "pd.DataFrame": pd.DataFrame(np.random.random(size=(1, 20))),
 }
 MULTIVARIATE = {
     "np.ndarray": np.random.random(size=(5, 20)),
-    "pd.Series": pd.Series(np.random.random(size=(20))),
     "pd.DataFrame": pd.DataFrame(np.random.random(size=(5, 20))),
 }
+
+UNIVARIATE_MISSING = {
+    "np.ndarray": UNIVARIATE["np.ndarray"].copy(),
+    "pd.Series": UNIVARIATE["pd.Series"].copy(),
+    "pd.DataFrame": UNIVARIATE["pd.DataFrame"].copy(),
+}
+UNIVARIATE_MISSING["np.ndarray"][np.random.randint(20)] = np.NAN
+UNIVARIATE_MISSING["np.ndarray"][np.random.randint(20)] = np.NAN
+UNIVARIATE_MISSING["pd.Series"][np.random.randint(20)] = np.NAN
+UNIVARIATE_MISSING["pd.Series"][np.random.randint(20)] = np.NAN
+UNIVARIATE_MISSING["pd.DataFrame"].iloc[0, np.random.randint(20)] = np.NAN
+UNIVARIATE_MISSING["pd.DataFrame"].iloc[0, np.random.randint(20)] = np.NAN
+
+MULTIVARIATE_MISSING = {
+    "np.ndarray": MULTIVARIATE["np.ndarray"].copy(),
+    "pd.DataFrame": MULTIVARIATE["pd.DataFrame"].copy(),
+}
+MULTIVARIATE_MISSING["np.ndarray"][np.random.randint(5)][np.random.randint(20)] = np.NAN
+MULTIVARIATE_MISSING["np.ndarray"][np.random.randint(5)][np.random.randint(20)] = np.NAN
+MULTIVARIATE_MISSING["pd.DataFrame"].iloc[
+    np.random.randint(5), np.random.randint(20)
+] = np.NAN
+MULTIVARIATE_MISSING["pd.DataFrame"].iloc[
+    np.random.randint(5), np.random.randint(20)
+] = np.NAN
+
+
 VALID_TYPES = [
     "np.ndarray",
     "pd.Series",
@@ -85,116 +50,364 @@ VALID_TYPES = [
 ]
 
 
-def test_univariate_convert_X():
-    """Test _convert_X on univariate series."""
-    # np.ndarray inner, no multivariate. Univariate are always series,
-    # not dataframes, axis is ignored
-    dummy1 = BaseSeriesEstimator()
-    for st in ["np.ndarray", "pd.Series"]:  # Inner type, DataFrame not allowed for
-        # univariate
-        for st2 in VALID_TYPES:  # input type
-            X = UNIVARIATE[st2]
-            dummy1.set_tags(**{"X_inner_type": st})
-            X2 = dummy1._convert_X(X, axis=0)
-            assert type(X2).__name__ == st.split(".")[1]
-            X2 = dummy1._convert_X(X, axis=1)
-            assert len(X) == len(X2)
-    with pytest.raises(ValueError, match="Axis should be 0 or 1"):
-        dummy1._convert_X(X, axis=2)
-    d = BaseSeriesEstimator.get_test_params()
-    assert d["axis"] == 0
-    all_tags = {"X_inner_type": "numpy3D"}
-    dummy1.set_tags(**all_tags)
-    with pytest.raises(ValueError, match="Unsupported inner"):
-        dummy1._convert_X(X, axis=0)
-
-
-def test_multivariate_convert_X():
-    """Test _convert_X on multivariate series."""
-    # np.ndarray inner multivariate capable. Multivariate are always DataFrames,
-    # not Series. 1D np.ndarray converted to 2D. Capability is tested in check_X,
-    # not here
-    dummy1 = BaseSeriesEstimator(axis=0)
-    for st in ["np.ndarray", "pd.DataFrame"]:  # Inner type, Series not allowed for
-        # Multivariate
-        name = st.split(".")[1]
-        for st2 in VALID_TYPES:  # input type
-            dummy1.set_tags(**{"X_inner_type": st, "capability:multivariate": True})
-            X = MULTIVARIATE[st2]
-            X2 = dummy1._convert_X(X, axis=0)
-            assert type(X2).__name__ == name
-
-            shape1 = X2.shape
-            X2 = dummy1._convert_X(X, axis=1)
-            shape2 = X2.shape
-            assert shape1[0] == shape2[1] and shape1[1] == shape2[0]
+def test_check_X():
+    """Test if capabilities correctly tested in _check_X."""
+    dummy = BaseSeriesEstimator(axis=1)
+
+    # check basic univariate input
+    meta = dummy._check_X(UNIVARIATE["np.ndarray"], axis=1)
+    meta2 = dummy._check_X(UNIVARIATE["pd.Series"], axis=1)
+    meta3 = dummy._check_X(UNIVARIATE["pd.DataFrame"], axis=1)
+    assert meta == meta2 == meta3
+    assert not meta["multivariate"]
+    assert meta["n_channels"] == 1
+    assert not meta["missing_values"]
+
+    # DataFrames are always 2d
+    meta = dummy._check_X(UNIVARIATE["pd.DataFrame"].T, axis=0)
+    assert not meta["multivariate"]
+    assert meta["n_channels"] == 1
+    assert not meta["missing_values"]
+
+    # check multivariate capability False
+    with pytest.raises(ValueError, match="Multivariate data not supported"):
+        dummy._check_X(MULTIVARIATE["np.ndarray"], axis=1)
+    with pytest.raises(ValueError, match="Multivariate data not supported"):
+        dummy._check_X(MULTIVARIATE["pd.DataFrame"], axis=1)
+    with pytest.raises(ValueError, match="Multivariate data not supported"):
+        dummy._check_X(MULTIVARIATE["np.ndarray"].T, axis=0)
+    with pytest.raises(ValueError, match="Multivariate data not supported"):
+        dummy._check_X(MULTIVARIATE["pd.DataFrame"].T, axis=0)
+
+    # check missing value capability False
+    dummy.set_tags(**{"capability:multivariate": True})
+
+    with pytest.raises(ValueError, match="Missing values not supported"):
+        dummy._check_X(UNIVARIATE_MISSING["np.ndarray"], axis=1)
+    with pytest.raises(ValueError, match="Missing values not supported"):
+        dummy._check_X(UNIVARIATE_MISSING["pd.Series"], axis=1)
+    with pytest.raises(ValueError, match="Missing values not supported"):
+        dummy._check_X(UNIVARIATE_MISSING["pd.DataFrame"], axis=1)
+    with pytest.raises(ValueError, match="Missing values not supported"):
+        dummy._check_X(MULTIVARIATE_MISSING["np.ndarray"], axis=1)
+    with pytest.raises(ValueError, match="Missing values not supported"):
+        dummy._check_X(MULTIVARIATE_MISSING["pd.DataFrame"], axis=1)
+
+    # check multivariate capable
+    meta = dummy._check_X(MULTIVARIATE["np.ndarray"], axis=1)
+    meta2 = dummy._check_X(MULTIVARIATE["pd.DataFrame"], axis=1)
+    meta3 = dummy._check_X(MULTIVARIATE["np.ndarray"].T, axis=0)
+    meta4 = dummy._check_X(MULTIVARIATE["pd.DataFrame"].T, axis=0)
+    assert meta == meta2 == meta3 == meta4
+    assert meta["multivariate"]
+    assert meta["n_channels"] == 5
+    assert not meta["missing_values"]
+
+    # check missing value capable
+    dummy.set_tags(**{"capability:missing_values": True})
+
+    meta = dummy._check_X(UNIVARIATE_MISSING["np.ndarray"], axis=1)
+    meta2 = dummy._check_X(UNIVARIATE_MISSING["pd.Series"], axis=1)
+    meta3 = dummy._check_X(UNIVARIATE_MISSING["pd.DataFrame"], axis=1)
+    meta4 = dummy._check_X(MULTIVARIATE_MISSING["np.ndarray"], axis=1)
+    meta5 = dummy._check_X(MULTIVARIATE_MISSING["pd.DataFrame"], axis=1)
+    assert meta == meta2 == meta3
+    assert meta4 == meta5
+    assert meta["missing_values"] and meta4["missing_values"]
+    assert not meta["multivariate"]
+    assert meta["n_channels"] == 1
+    assert meta4["multivariate"]
+    assert meta4["n_channels"] == 5
+
+    # check univariate capability False
+    dummy.set_tags(**{"capability:univariate": False})
+
+    with pytest.raises(ValueError, match="Univariate data not supported"):
+        dummy._check_X(UNIVARIATE["np.ndarray"], axis=1)
+    with pytest.raises(ValueError, match="Univariate data not supported"):
+        dummy._check_X(UNIVARIATE["pd.Series"], axis=1)
+    with pytest.raises(ValueError, match="Univariate data not supported"):
+        dummy._check_X(UNIVARIATE["pd.DataFrame"], axis=1)
+    with pytest.raises(ValueError, match="Univariate data not supported"):
+        dummy._check_X(UNIVARIATE["pd.DataFrame"].T, axis=0)
+
+    # check invalid inputs and data types
+    with pytest.raises(ValueError, match="Input type of X should be one of"):
+        dummy._check_X("String input", axis=1)
+
+    invalid_always = np.array(["1", "2", "3", "4", "5"])
+    with pytest.raises(ValueError, match="dtype for np.ndarray must be float or int"):
+        dummy._check_X(invalid_always, axis=1)
+    with pytest.raises(ValueError, match="pd.Series dtype must be numeric"):
+        dummy._check_X(pd.Series(invalid_always), axis=1)
+    with pytest.raises(ValueError, match="pd.DataFrame dtype must be numeric"):
+        dummy._check_X(pd.DataFrame(invalid_always), axis=1)
+
+    with pytest.raises(ValueError, match="Input axis should be 0 or 1, saw 2"):
+        dummy._check_X(UNIVARIATE["np.ndarray"], axis=2)
+    with pytest.raises(ValueError, match="Input axis should be 0 or 1, saw -1"):
+        dummy._check_X(UNIVARIATE["np.ndarray"], axis=-1)
+
+    collection = np.random.random(size=(10, 2, 10))
+    with pytest.raises(ValueError, match="X must have at most 2 dimensions"):
+        dummy._check_X(collection, axis=0)
+
+
+@pytest.mark.parametrize("input_type", VALID_TYPES)
+def test_convert_X_ndarray_inner(input_type):
+    """Test _convert_X on with np.ndarray inner type."""
+    dummy = BaseSeriesEstimator(axis=1)
+    dummy.set_tags(**{"X_inner_type": "np.ndarray"})
+    # test univariate
+    X = UNIVARIATE[input_type]
+
+    X2 = dummy._convert_X(X, axis=1)
+    assert type(X2).__name__ == "ndarray"
+    assert X2.shape == (1, 20)
+    assert X.shape[-1] == X2.shape[1]
+
+    X3 = dummy._convert_X(X.T, axis=0)
+    assert type(X3).__name__ == "ndarray"
+    assert X3.shape == (1, 20)
+    assert_equal(X2, X3)
+
+    # test multivariate
+    if input_type != "pd.Series":
+        dummy.set_tags(**{"capability:multivariate": True})
+        X = MULTIVARIATE[input_type]
+
+        X2 = dummy._convert_X(X, axis=1)
+        assert type(X2).__name__ == "ndarray"
+        assert X2.shape == (5, 20)
+        assert X.shape == X2.shape
+
+        X3 = dummy._convert_X(X.T, axis=0)
+        assert type(X3).__name__ == "ndarray"
+        assert X3.shape == (5, 20)
+        assert_equal(X2, X3)
+
+
+@pytest.mark.parametrize("input_type", VALID_TYPES)
+def test_convert_X_series_inner(input_type):
+    """Test _convert_X on with pd.Series inner type."""
+    dummy = BaseSeriesEstimator(axis=1)
+    dummy.set_tags(**{"X_inner_type": "pd.Series"})
+    # test univariate
+    X = UNIVARIATE[input_type]
+
+    X2 = dummy._convert_X(X, axis=1)
+    assert type(X2).__name__ == "Series"
+    assert X2.shape == (20,)
+    assert X.shape[-1] == X2.shape[0]
+
+    X3 = dummy._convert_X(X.T, axis=0)
+    assert type(X3).__name__ == "Series"
+    assert X3.shape == (20,)
+    assert X2.equals(X3)
+
+    # test multivariate, should raise error as pd.Series cannot be multivariate
+    if input_type != "pd.Series":
+        X = MULTIVARIATE[input_type]
+
+        with pytest.raises(ValueError, match="multivariate data. Found 5 channels"):
+            dummy._convert_X(X, axis=1)
+
+        dummy.set_tags(**{"capability:multivariate": True})
+
+        with pytest.raises(ValueError, match="for multivariate capable estimators"):
+            dummy._convert_X(X, axis=1)
+        with pytest.raises(ValueError, match="for multivariate capable estimators"):
+            dummy._convert_X(X.T, axis=0)
+
+
+@pytest.mark.parametrize("input_type", VALID_TYPES)
+def test_convert_X_dataframe_inner(input_type):
+    """Test _convert_X on with pd.DataFrame inner type."""
+    dummy = BaseSeriesEstimator(axis=1)
+    dummy.set_tags(**{"X_inner_type": "pd.DataFrame"})
+    # test univariate
+    X = UNIVARIATE[input_type]
+
+    X2 = dummy._convert_X(X, axis=1)
+    assert type(X2).__name__ == "DataFrame"
+    assert X2.shape == (1, 20)
+    assert X.shape[-1] == X2.shape[1]
+
+    X3 = dummy._convert_X(X.T, axis=0)
+    assert type(X3).__name__ == "DataFrame"
+    assert X3.shape == (1, 20)
+    assert X2.equals(X3)
+
+    # test multivariate
+    if input_type != "pd.Series":
+        dummy.set_tags(**{"capability:multivariate": True})
+        X = MULTIVARIATE[input_type]
+
+        X2 = dummy._convert_X(X, axis=1)
+        assert type(X2).__name__ == "DataFrame"
+        assert X2.shape == (5, 20)
+        assert X.shape == X2.shape
+
+        X3 = dummy._convert_X(X.T, axis=0)
+        assert type(X3).__name__ == "DataFrame"
+        assert X3.shape == (5, 20)
+        assert X2.equals(X3)
+
+
+def test_convert_X_invalid():
+    """Test _convert_X for invalid inputs."""
+    dummy = BaseSeriesEstimator(axis=1)
+
+    with pytest.raises(ValueError, match="Input axis should be 0 or 1, saw 2"):
+        dummy._convert_X(UNIVARIATE["np.ndarray"], axis=2)
+    with pytest.raises(ValueError, match="Input axis should be 0 or 1, saw -1"):
+        dummy._convert_X(UNIVARIATE["np.ndarray"], axis=-1)
+
+    dummy.set_tags(**{"X_inner_type": ["invalid"]})
+    with pytest.raises(ValueError, match="Unsupported inner type"):
+        dummy._convert_X(UNIVARIATE["np.ndarray"], axis=1)
 
 
 @pytest.mark.parametrize("input_type", VALID_TYPES)
 @pytest.mark.parametrize("inner_type", VALID_TYPES)
 def test_preprocess_series(input_type, inner_type):
-    """Test _preprocess_series."""
-    dummy1 = BaseSeriesEstimator(axis=0)
-    inner = inner_type.split(".")[1]
-    if inner_type != "pd.DataFrame":
-        dummy1.set_tags(**{"X_inner_type": inner_type})
-        X = UNIVARIATE[input_type]
-        X2 = dummy1._preprocess_series(X)
-        assert type(X2).__name__ == inner
-    if inner_type != "pd.Series":
-        dummy1.set_tags(**{"X_inner_type": inner_type, "capability:multivariate": True})
-        X = UNIVARIATE[input_type]
-        X2 = dummy1._preprocess_series(X)
-        assert type(X2).__name__ == inner
-
-
-INPUT_CORRECT = [
-    np.array([0, 0, 0, 1, 1]),
-    pd.Series([0, 0, 0, 1, 1]),
-    pd.DataFrame([0, 0, 0, 1, 1]),
-]
-INPUT_WRONG = [
-    np.array([[0, 0, 0, 1, 1, 2], [0, 0, 0, 1, 1, 2]]),
-    pd.DataFrame([[0, 0, 0, 1, 1, 2], [0, 0, 0, 1, 1, 2]]),
-    np.array([0, 0, 0, 1, "FOO"]),
-    pd.Series([0, 0, 0, 1, "FOO"]),
-    pd.DataFrame([0, 0, 0, 1, "FOO"]),
-    "Up the arsenal",
-]
+    """Test _preprocess_series for different input and inner types."""
+    dummy = BaseSeriesEstimator(axis=1)
+    dummy.set_tags(**{"X_inner_type": inner_type})
+    inner_name = inner_type.split(".")[1]
+    # test univariate
+    X = UNIVARIATE[input_type]
 
+    X2 = dummy._preprocess_series(X, axis=1, store_metadata=True)
+    assert type(X2).__name__ == inner_name
+    assert X.shape[-1] == X2.shape[-1]
+    assert not dummy.metadata_["multivariate"]
+    assert dummy.metadata_["n_channels"] == 1
+    assert not dummy.metadata_["missing_values"]
 
-@pytest.mark.parametrize("y_correct", INPUT_CORRECT)
-def test_check_y_correct(y_correct):
-    """Test the _check_y method with correct input."""
-    assert BaseSeriesEstimator._check_y(None, y_correct) is None
+    X3 = dummy._preprocess_series(X.T, axis=0, store_metadata=True)
+    assert type(X3).__name__ == inner_name
+    assert X.shape[-1] == X3.shape[-1]
+    assert not dummy.metadata_["multivariate"]
+    assert dummy.metadata_["n_channels"] == 1
+    assert not dummy.metadata_["missing_values"]
 
+    # test multivariate, excludes pd.Series input type as it cannot be multivariate
+    if input_type != "pd.Series" and inner_type != "pd.Series":
+        dummy.set_tags(**{"capability:multivariate": True})
+        X = MULTIVARIATE[input_type]
 
-@pytest.mark.parametrize("y_wrong", INPUT_WRONG)
-def test_check_y_wrong(y_wrong):
-    """Test the _check_y method with incorrect input."""
-    with pytest.raises(ValueError, match="Error in input type for y"):
-        BaseSeriesEstimator._check_y(None, y_wrong)
+        X2 = dummy._preprocess_series(X, axis=1, store_metadata=True)
+        assert type(X2).__name__ == inner_name
+        assert X2.shape == (5, 20)
+        assert X.shape == X2.shape
+        assert dummy.metadata_["multivariate"]
+        assert dummy.metadata_["n_channels"] == 5
+        assert not dummy.metadata_["missing_values"]
+
+        X3 = dummy._preprocess_series(X.T, axis=0, store_metadata=True)
+        assert type(X3).__name__ == inner_name
+        assert X3.shape == (5, 20)
+        assert X.shape == X3.shape
+        assert dummy.metadata_["multivariate"]
+        assert dummy.metadata_["n_channels"] == 5
+        assert not dummy.metadata_["missing_values"]
+    elif input_type != "pd.Series":
+        dummy.set_tags(**{"capability:multivariate": True})
+        X = MULTIVARIATE[input_type]
+
+        with pytest.raises(ValueError, match="Cannot convert to pd.Series"):
+            dummy._preprocess_series(X, axis=1, store_metadata=True)
 
 
 @pytest.mark.parametrize("returned_type", VALID_TYPES)
-def test__postprocess_series(returned_type):
-    """Test data reformatted correctly."""
-    x = UNIVARIATE[returned_type]
-    dummy1 = BaseSeriesEstimator()
-    tags = {"capability:multivariate": False}
-    dummy1.set_tags(**tags)
-    y = dummy1._postprocess_series(x)
-    y2 = dummy1._postprocess_series(x, axis=0)
-    y3 = dummy1._postprocess_series(x, axis=1)
-    assert y.ndim == 1
-    assert y2.shape == y3.shape
-    tags = {"capability:multivariate": True}
-    dummy1.set_tags(**tags)
-    x = MULTIVARIATE[returned_type]
-    y = dummy1._postprocess_series(x)
-    y2 = dummy1._postprocess_series(x, axis=0)
-    y3 = dummy1._postprocess_series(x, axis=1)
-    assert x.shape == y.shape
-    assert y.shape == y2.shape
+def test_postprocess_series(returned_type):
+    """Test data reformatted correctly by _postprocess_series."""
+    dummy = BaseSeriesEstimator(axis=1)
+    # test univariate
+    Xu = UNIVARIATE[returned_type]
+
+    X2 = dummy._postprocess_series(Xu, axis=1)
+    assert X2.ndim == 1
+    assert X2.shape == (20,)
+
+    X3 = dummy._postprocess_series(Xu.T, axis=0)
+    assert X2.shape == X3.shape
+
+    # test multivariate and 2d, excludes pd.Series input type as it is 1D only
     if returned_type != "pd.Series":
-        assert y.shape[0] == y3.shape[1] and y.shape[1] == y3.shape[0]
+        Xu = Xu.reshape((1, -1)) if returned_type == "np.ndarray" else Xu
+        Xm = MULTIVARIATE[returned_type]
+
+        X2 = dummy._postprocess_series(Xu, axis=1)
+        assert X2.ndim == 1
+        assert X2.shape == (20,)
+
+        X3 = dummy._postprocess_series(Xu.T, axis=0)
+        assert X2.shape == X3.shape
+
+        X2 = dummy._postprocess_series(Xm, axis=1)
+        assert X2.ndim == 2
+        assert X2.shape == (5, 20)
+
+        X3 = dummy._postprocess_series(Xm.T, axis=0)
+        assert X2.shape == X3.shape
+
+        # multivariate capable estimator
+        dummy.set_tags(**{"capability:multivariate": True})
+
+        X2 = dummy._postprocess_series(Xu, axis=1)
+        assert X2.ndim == 2
+        assert X2.shape == (1, 20)
+
+        X3 = dummy._postprocess_series(Xu.T, axis=0)
+        assert X2.shape == X3.shape
+
+        X2 = dummy._postprocess_series(Xm, axis=1)
+        assert X2.ndim == 2
+        assert X2.shape == (5, 20)
+
+        X3 = dummy._postprocess_series(Xm.T, axis=0)
+        assert X2.shape == X3.shape
+
+
+def test_axis():
+    """Test axis property."""
+    dummy = BaseSeriesEstimator(axis=1)
+    dummy.set_tags(**{"capability:multivariate": True})
+    X = MULTIVARIATE["np.ndarray"]
+
+    X2 = dummy._preprocess_series(X, axis=1, store_metadata=True)
+    assert type(X2).__name__ == "ndarray"
+    assert X2.shape == (5, 20)
+    assert X.shape == X2.shape
+    assert dummy.metadata_["multivariate"]
+    assert dummy.metadata_["n_channels"] == 5
+    assert not dummy.metadata_["missing_values"]
+
+    X2 = dummy._preprocess_series(X.T, axis=0, store_metadata=True)
+    assert type(X2).__name__ == "ndarray"
+    assert X2.shape == (5, 20)
+    assert X.shape == X2.shape
+    assert dummy.metadata_["multivariate"]
+    assert dummy.metadata_["n_channels"] == 5
+    assert not dummy.metadata_["missing_values"]
+
+    dummy.axis = 0
+    X = X.T
+
+    X2 = dummy._preprocess_series(X.T, axis=1, store_metadata=True)
+    assert type(X2).__name__ == "ndarray"
+    assert X2.shape == (20, 5)
+    assert X.shape == X2.shape
+    assert dummy.metadata_["multivariate"]
+    assert dummy.metadata_["n_channels"] == 5
+    assert not dummy.metadata_["missing_values"]
+
+    X2 = dummy._preprocess_series(X, axis=0, store_metadata=True)
+    assert type(X2).__name__ == "ndarray"
+    assert X2.shape == (20, 5)
+    assert X.shape == X2.shape
+    assert dummy.metadata_["multivariate"]
+    assert dummy.metadata_["n_channels"] == 5
+    assert not dummy.metadata_["missing_values"]
diff --git a/aeon/clustering/averaging/tests/test_dba.py b/aeon/clustering/averaging/tests/test_dba.py
index 7836c8b1..dd01f48f 100644
--- a/aeon/clustering/averaging/tests/test_dba.py
+++ b/aeon/clustering/averaging/tests/test_dba.py
@@ -1,146 +1,193 @@
 """Tests for DBA."""
 
+import re
+
 import numpy as np
 import pytest
 
-from aeon.clustering.averaging import elastic_barycenter_average
-from aeon.testing.utils.data_gen import make_example_3d_numpy
+from aeon.clustering.averaging import (
+    elastic_barycenter_average,
+    petitjean_barycenter_average,
+    subgradient_barycenter_average,
+)
+from aeon.testing.utils.data_gen import make_example_3d_numpy, make_series
 
-expected_dba = np.array(
+expected_petitjean_dba_univariate = np.array(
     [
         [
-            2.14365783,
-            1.79294991,
-            1.40583643,
-            1.41551047,
-            1.35200733,
-            1.89139937,
-            1.86542327,
-            1.32466944,
-            1.67397517,
-            1.8709378,
-        ],
-        [
-            1.39792187,
-            1.48779949,
-            1.26747876,
-            1.77152767,
-            1.89443306,
-            1.33528121,
-            1.8849385,
-            1.61457405,
-            1.6470203,
-            1.58740109,
-        ],
-        [
-            1.20289266,
-            0.99392693,
-            0.91151178,
-            1.36251883,
-            1.63591345,
-            1.76292241,
-            1.2924362,
-            1.82067719,
-            1.36224209,
-            2.09249394,
-        ],
+            1.4108696053220873,
+            2.394226235189873,
+            1.1060488942437015,
+            2.623297561062815,
+            0.527766384447058,
+            1.43188098329184,
+            1.6236753043608105,
+            1.2058072160220672,
+            1.1967764056706969,
+            2.0004477079532927,
+        ]
+    ]
+)
+
+expected_petitjean_dba_multivariate = np.array(
+    [
         [
-            1.50362661,
-            1.33955338,
-            2.21909116,
-            1.84855474,
-            1.31602877,
-            1.53317511,
-            1.84644465,
-            1.42878873,
-            1.78336225,
-            1.59598453,
+            0.7267915954182468,
+            1.3623320157069627,
+            1.4495279596147028,
+            0.5926686136095334,
+            0.9971756690146295,
+            1.3283383084721943,
+            0.7659577046471742,
+            1.9660527002664323,
+            1.063255339395868,
+            1.58993866039882,
         ],
         [
-            1.76544264,
-            1.41121261,
-            1.63573482,
-            1.02127341,
-            1.41004018,
-            1.70329872,
-            1.58462695,
-            1.59081868,
-            1.44386256,
-            1.07904533,
+            1.2617459283113774,
+            1.205815463685702,
+            1.2234939638872127,
+            1.7840697833529937,
+            0.6722624075103683,
+            1.835414289773766,
+            0.9514286846149641,
+            1.2908899681880017,
+            1.416361929410518,
+            1.5241706749949016,
         ],
         [
-            1.30443531,
-            1.58311813,
-            1.70761705,
-            2.20491277,
-            1.26669821,
-            1.70208611,
-            1.91908395,
-            1.60117971,
-            1.90385381,
-            1.64541667,
+            1.3034322787744397,
+            1.5436844012461954,
+            1.632959525575443,
+            1.0384357523605505,
+            1.250310558432616,
+            0.7533753079264169,
+            0.7766566039307113,
+            1.2180892807545585,
+            1.6222668068442372,
+            1.3287035370261573,
         ],
+    ]
+)
+
+expected_subgradient_dba_univariate = np.array(
+    [
         [
-            1.20010098,
-            2.06998292,
-            1.75802436,
-            1.55995949,
-            1.8326049,
-            1.05988814,
-            2.31728009,
-            1.47100396,
-            2.19823014,
-            1.19986879,
-        ],
+            1.51074199,
+            2.15864417,
+            1.00234603,
+            2.35561396,
+            0.63780835,
+            1.40458666,
+            1.51369944,
+            1.2566171,
+            1.18067802,
+            1.8077327,
+        ]
+    ]
+)
+
+expected_subgradient_dba_multivariate = np.array(
+    [
         [
-            1.55797104,
-            1.40613571,
-            1.85809765,
-            1.41017291,
-            1.9308945,
-            1.37990633,
-            1.77000225,
-            1.35118687,
-            1.25940145,
-            1.85597143,
+            0.74762929,
+            1.27990802,
+            1.39635383,
+            0.77782365,
+            0.96456039,
+            1.33559105,
+            0.71501953,
+            1.68703826,
+            1.03765683,
+            1.59958969,
         ],
         [
-            1.85856132,
-            2.09754471,
-            1.47383472,
-            1.51400218,
-            1.43779312,
-            1.69484847,
-            1.95836839,
-            1.62152772,
-            1.26513914,
-            2.08603391,
+            1.28770927,
+            1.23770584,
+            1.20206139,
+            1.78964445,
+            0.89355146,
+            1.6961715,
+            1.00333335,
+            1.23876289,
+            1.38498497,
+            1.55936623,
         ],
         [
-            1.2279546,
-            1.53877513,
-            1.32957377,
-            1.71607611,
-            1.46982699,
-            1.95766663,
-            1.4674824,
-            1.85898204,
-            1.81942838,
-            1.6760823,
+            1.34461418,
+            1.50480679,
+            1.55114792,
+            1.18503368,
+            1.04103323,
+            0.90804179,
+            1.04520108,
+            1.07537586,
+            1.57586214,
+            1.34801647,
         ],
     ]
 )
 
 
-def test_dba():
-    """Test dba functionality."""
-    X_train = make_example_3d_numpy(10, 10, 10, random_state=1, return_y=False)
+def test_petitjean_dba():
+    """Test petitjean dba functionality."""
+    X_train_uni = make_example_3d_numpy(10, 1, 10, random_state=1, return_y=False)
 
-    average_ts = elastic_barycenter_average(X_train)
+    average_ts_uni = elastic_barycenter_average(
+        X_train_uni, method="petitjean", random_state=1
+    )
+    call_directly_average_ts_uni = petitjean_barycenter_average(
+        X_train_uni, random_state=1
+    )
+    assert isinstance(average_ts_uni, np.ndarray)
+    assert average_ts_uni.shape == X_train_uni[0].shape
+    assert np.allclose(average_ts_uni, expected_petitjean_dba_univariate)
+    assert np.allclose(average_ts_uni, call_directly_average_ts_uni)
 
-    assert isinstance(average_ts, np.ndarray)
-    assert average_ts.shape == X_train[0].shape
-    assert np.allclose(average_ts, expected_dba)
+    X_train_multi = make_example_3d_numpy(10, 3, 10, random_state=1, return_y=False)
+
+    average_ts_multi = elastic_barycenter_average(
+        X_train_multi, method="petitjean", random_state=1
+    )
+    call_directly_average_ts_multi = petitjean_barycenter_average(
+        X_train_multi, random_state=1
+    )
+
+    assert isinstance(average_ts_multi, np.ndarray)
+    assert average_ts_multi.shape == X_train_multi[0].shape
+    assert np.allclose(average_ts_multi, expected_petitjean_dba_multivariate)
+    assert np.allclose(average_ts_multi, call_directly_average_ts_multi)
+
+
+def test_subgradient_dba():
+    """Test stochastic subgradient dba functionality."""
+    X_train_uni = make_example_3d_numpy(10, 1, 10, random_state=1, return_y=False)
+
+    average_ts_uni = elastic_barycenter_average(
+        X_train_uni, method="subgradient", random_state=1
+    )
+    call_directly_average_ts_uni = subgradient_barycenter_average(
+        X_train_uni, random_state=1
+    )
+
+    assert isinstance(average_ts_uni, np.ndarray)
+    assert average_ts_uni.shape == X_train_uni[0].shape
+    assert np.allclose(average_ts_uni, expected_subgradient_dba_univariate)
+    assert np.allclose(average_ts_uni, call_directly_average_ts_uni)
+
+    X_train_multi = make_example_3d_numpy(10, 3, 10, random_state=1, return_y=False)
+
+    average_ts_multi = elastic_barycenter_average(
+        X_train_multi, method="subgradient", random_state=1
+    )
+    call_directly_average_ts_multi = subgradient_barycenter_average(
+        X_train_multi, random_state=1
+    )
+
+    assert isinstance(average_ts_multi, np.ndarray)
+    assert average_ts_multi.shape == X_train_multi[0].shape
+    assert np.allclose(average_ts_multi, expected_subgradient_dba_multivariate)
+    assert np.allclose(average_ts_multi, call_directly_average_ts_multi)
 
 
 @pytest.mark.parametrize(
@@ -166,5 +213,103 @@ def test_elastic_dba_variations(distance):
         X_train, distance=distance, window=0.2, independent=False
     )
 
+    subgradient_ts = elastic_barycenter_average(
+        X_train, distance=distance, window=0.2, independent=False, method="subgradient"
+    )
+
     assert isinstance(average_ts, np.ndarray)
+    assert isinstance(subgradient_ts, np.ndarray)
     assert average_ts.shape == X_train[0].shape
+    assert subgradient_ts.shape == X_train[0].shape
+
+
+@pytest.mark.parametrize(
+    "init_barycenter",
+    [
+        "mean",
+        "medoids",
+        "random",
+        (
+            make_series(10, 1, return_numpy=True, random_state=1),
+            make_series(
+                n_timepoints=4, n_columns=10, return_numpy=True, random_state=1
+            ),
+        ),
+    ],
+)
+def test_dba_init(init_barycenter):
+    """Test dba init functionality."""
+    if isinstance(init_barycenter, str):
+        univariate_init = init_barycenter
+        multivariate_init = init_barycenter
+    else:
+        univariate_init = init_barycenter[0]
+        multivariate_init = init_barycenter[1]
+
+    X_train_uni = make_example_3d_numpy(4, 1, 10, random_state=1, return_y=False)
+    average_ts_univ = elastic_barycenter_average(
+        X_train_uni, window=0.2, init_barycenter=univariate_init
+    )
+    subgradient_ts_univ = elastic_barycenter_average(
+        X_train_uni, window=0.2, method="subgradient", init_barycenter=univariate_init
+    )
+
+    assert isinstance(average_ts_univ, np.ndarray)
+    assert isinstance(subgradient_ts_univ, np.ndarray)
+
+    X_train_multi = make_example_3d_numpy(4, 4, 10, random_state=1, return_y=False)
+    average_ts_multi = elastic_barycenter_average(
+        X_train_multi, window=0.2, init_barycenter=multivariate_init
+    )
+    subgradient_ts_multi = elastic_barycenter_average(
+        X_train_multi,
+        window=0.2,
+        method="subgradient",
+        init_barycenter=multivariate_init,
+    )
+
+    assert isinstance(average_ts_multi, np.ndarray)
+    assert isinstance(subgradient_ts_multi, np.ndarray)
+
+
+def test_incorrect_input():
+    """Test dba incorrect input."""
+    # Test invalid distance
+    X = make_example_3d_numpy(10, 1, 10, return_y=False)
+    with pytest.raises(ValueError, match="Distance parameter invalid"):
+        elastic_barycenter_average(X, distance="Distance parameter invalid")
+
+    # Test invalid init barycenter string
+    with pytest.raises(
+        ValueError,
+        match="init_barycenter string is invalid. Please use one of the "
+        "following: 'mean', 'medoids', 'random'",
+    ):
+        elastic_barycenter_average(X, init_barycenter="init parameter invalid")
+
+    # Test invalid init barycenter type
+    with pytest.raises(
+        ValueError,
+        match="init_barycenter parameter is invalid. It must either be a "
+        "str or a np.ndarray",
+    ):
+        elastic_barycenter_average(X, init_barycenter=[[1, 2, 3]])
+
+    # Test invalid init barycenter with wrong shape
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "init_barycenter shape is invalid. Expected (1, 10) but " "got (1, 9)"
+        ),
+    ):
+        elastic_barycenter_average(X, init_barycenter=make_series(9, return_numpy=True))
+
+    # Test invalid berycenter method
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Invalid method: Not a real method. Please use one of the "
+            "following: ['petitjean', 'subgradient']"
+        ),
+    ):
+        elastic_barycenter_average(X, method="Not a real method")
diff --git a/aeon/clustering/tests/test_k_means.py b/aeon/clustering/tests/test_k_means.py
index 7bfc6f40..7baa98ad 100644
--- a/aeon/clustering/tests/test_k_means.py
+++ b/aeon/clustering/tests/test_k_means.py
@@ -247,6 +247,23 @@ def test_different_ba():
     # Check not equal to default
     assert not np.array_equal(default_mba, mba_custom_params_window_and_indep)
 
+    # Test subgradient ba
+    subgradient_dba_specified = _get_model_centres(
+        data,
+        distance="dtw",
+        average_params={"distance": "dtw", "method": "subgradient"},
+    )
+    assert not np.array_equal(dba_specified, subgradient_dba_specified)
+
+    # Test another distance and check passing custom distance param
+    subgradient_mba_specified = _get_model_centres(
+        data,
+        distance="msm",
+        average_params={"distance": "msm", "method": "subgradient"},
+    )
+    assert not np.array_equal(mba_specified, subgradient_mba_specified)
+    assert not np.array_equal(subgradient_dba_specified, subgradient_mba_specified)
+
 
 def check_value_in_every_cluster(num_clusters, initial_centres):
     """Check that every cluster has at least one value."""
@@ -311,8 +328,16 @@ def test_custom_distance_params():
     data = X_train[0:num_test_values]
 
     # Test passing distance param
-    default_dist = _get_model_centres(data, distance="msm")
+    default_dist = _get_model_centres(
+        data,
+        distance="msm",
+        distance_params={"window": 0.2},
+        average_params={"init_barycenter": "medoids"},
+    )
     custom_params_dist = _get_model_centres(
-        data, distance="msm", distance_params={"window": 0.2}
+        data,
+        distance="msm",
+        distance_params={"window": 0.2},
+        average_params={"init_barycenter": "mean"},
     )
     assert not np.array_equal(default_dist, custom_params_dist)
diff --git a/aeon/segmentation/tests/test_base.py b/aeon/segmentation/tests/test_base.py
index 14a7c501..013e4ecd 100644
--- a/aeon/segmentation/tests/test_base.py
+++ b/aeon/segmentation/tests/test_base.py
@@ -1,6 +1,7 @@
 """Test base segmenter."""
 
 import numpy as np
+import pandas as pd
 import pytest
 
 from aeon.segmentation.base import BaseSegmenter
@@ -53,3 +54,31 @@ def test_to_clusters():
     assert np.array_equal(labels, np.array([0, 0, 1, 1, 1, 1, 1, 1, 2, 2]))
     labels = BaseSegmenter.to_clusters([1, 2, 3], 4)
     assert np.array_equal(labels, np.array([0, 1, 2, 3]))
+
+
+INPUT_CORRECT = [
+    np.array([0, 0, 0, 1, 1]),
+    pd.Series([0, 0, 0, 1, 1]),
+    pd.DataFrame([0, 0, 0, 1, 1]),
+]
+INPUT_WRONG = [
+    np.array([[0, 0, 0, 1, 1, 2], [0, 0, 0, 1, 1, 2]]),
+    pd.DataFrame([[0, 0, 0, 1, 1, 2], [0, 0, 0, 1, 1, 2]]),
+    np.array([0, 0, 0, 1, "FOO"]),
+    pd.Series([0, 0, 0, 1, "FOO"]),
+    pd.DataFrame([0, 0, 0, 1, "FOO"]),
+    "Up the arsenal",
+]
+
+
+@pytest.mark.parametrize("y_correct", INPUT_CORRECT)
+def test_check_y_correct(y_correct):
+    """Test the _check_y method with correct input."""
+    assert BaseSegmenter._check_y(None, y_correct) is None
+
+
+@pytest.mark.parametrize("y_wrong", INPUT_WRONG)
+def test_check_y_wrong(y_wrong):
+    """Test the _check_y method with incorrect input."""
+    with pytest.raises(ValueError, match="Error in input type for y"):
+        BaseSegmenter._check_y(None, y_wrong)