Repository: raphaelvallat/yasa

Issue:
Hypnogram.hypno as a pd.Categorical?
@remrama , related to our discussion about the output of `yasa.Hypnogram.upsample_to_data` in #124, I was thinking that we could in fact set the default dtype of `hyp.hypno` as a [pandas.Categorical](https://pandas.pydata.org/pandas-docs/version/1.4/reference/api/pandas.Categorical.html). This reduces the memory size of an upsampled hypnogram (5 stages, 8 hour TIB, freq="10ms") from 22 MB (dtype=str or dtype=np.int64) to 2.7 MB.

```python
pd.Categorical(hyp.hypno, categories=hyp.labels, ordered=False)
```

Discussion:
codecov-commenter: # [Codecov](https://codecov.io/gh/raphaelvallat/yasa/pull/126?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Raphael+Vallat) Report
Base: **92.57**% // Head: **92.59**% // Increases project coverage by **`+0.01%`** :tada:
> Coverage data is based on head [(`64a9461`)](https://codecov.io/gh/raphaelvallat/yasa/pull/126?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Raphael+Vallat) compared to base [(`9be559d`)](https://codecov.io/gh/raphaelvallat/yasa/commit/9be559d23639d9cc59f1a23d4eb7fd366f4bb32b?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Raphael+Vallat).
> Patch coverage: 100.00% of modified lines in pull request are covered.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master     #126      +/-   ##
==========================================
+ Coverage   92.57%   92.59%   +0.01%     
==========================================
  Files          23       23              
  Lines        3112     3119       +7     
==========================================
+ Hits         2881     2888       +7     
  Misses        231      231              
```


| [Impacted Files](https://codecov.io/gh/raphaelvallat/yasa/pull/126?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Raphael+Vallat) | Coverage Δ | |
|---|---|---|
| [yasa/hypno.py](https://codecov.io/gh/raphaelvallat/yasa/pull/126/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Raphael+Vallat#diff-eWFzYS9oeXBuby5weQ==) | `98.58% <100.00%> (+0.01%)` | :arrow_up: |
| [yasa/tests/test\_hypnoclass.py](https://codecov.io/gh/raphaelvallat/yasa/pull/126/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Raphael+Vallat#diff-eWFzYS90ZXN0cy90ZXN0X2h5cG5vY2xhc3MucHk=) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Raphael+Vallat). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Raphael+Vallat)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/raphaelvallat/yasa/pull/126?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Raphael+Vallat).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Raphael+Vallat).


Patch:
diff --git a/docs/changelog.rst b/docs/changelog.rst
index 6bb1bcb..e4fb211 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -11,12 +11,12 @@ What's new
 v0.7.dev
 --------
 
-**Object-oriented hypnogram manipulation**
+**Object-oriented hypnogram**
 
-This version introduces the new :py:class:`yasa.Hypnogram` class, which will progressively become
+This version introduces the new :py:class:`yasa.Hypnogram` class, which will gradually become
 the standard way to store and manipulate hypnograms in YASA. Put simply, YASA now uses an
 object-oriented approach to hypnograms. That is, hypnograms are now stored as a class (aka object),
-which comes with several pre-built functions (named methods) and attributes. See for example below:
+which comes with several pre-built functions (aka methods) and attributes. See for example below:
 
 .. code-block::  python
 
@@ -25,7 +25,7 @@ which comes with several pre-built functions (named methods) and attributes. See
     values = ["W", "W", "W", "S", "S", "S", "S", "S", "W", "S", "S", "S"]
     hyp = Hypnogram(values, n_stages=2, start="2022-12-23 22:30:00", scorer="RM")
     # Below are some class attributes
-    hyp.hypno  # Hypnogram values (pandas Series)
+    hyp.hypno  # Hypnogram values (a pandas.Series of categorical dtype)
     hyp.duration  # Total duration of the hypnogram, in minutes
     hyp.sampling_frequency  # Sampling frequency of the hypnogram
     hyp.mapping  # Mapping from strings to integers
diff --git a/yasa/hypno.py b/yasa/hypno.py
index f3d6f92..81b0c80 100644
--- a/yasa/hypno.py
+++ b/yasa/hypno.py
@@ -9,6 +9,7 @@ import pandas as pd
 from yasa.io import set_log_level
 from yasa.plotting import plot_hypnogram
 from yasa.sleepstats import transition_matrix
+from pandas.api.types import CategoricalDtype
 
 __all__ = [
     "Hypnogram",
@@ -18,7 +19,7 @@ __all__ = [
     "hypno_upsample_to_data",
     "hypno_find_periods",
     "load_profusion_hypno",
-    "simulate_hypno",
+    "simulate_hypnogram",
 ]
 
 
@@ -29,13 +30,13 @@ class Hypnogram:
     """
     Main class for manipulating hypnogram in YASA.
 
-    Starting with YASA v0.7, YASA takes a more object-oriented approach to hypnograms. That is,
+    Starting with v0.7, YASA takes a more object-oriented approach to hypnograms. That is,
     hypnograms are now stored as a class (aka object), which comes with its own attributes and
     functions. Furthermore, YASA does not allow integer values to define the stages anymore.
     Instead, users must pass an array of strings with the actual stage names
     (e.g. ["WAKE", "WAKE", "N1", ..., "REM", "REM"]).
 
-    .. seealso:: :py:func:`yasa.simulate_hypno`
+    .. seealso:: :py:func:`yasa.simulate_hypnogram`
 
     .. versionadded:: 0.7.0
 
@@ -96,7 +97,8 @@ class Hypnogram:
     9     SLEEP
     10    SLEEP
     11    SLEEP
-    Name: Stage, dtype: object
+    Name: Stage, dtype: category
+    Categories (4, object): ['WAKE', 'SLEEP', 'ART', 'UNS']
 
     >>> # Number of epochs in the hypnogram
     >>> hyp.n_epochs
@@ -125,7 +127,7 @@ class Hypnogram:
     9     1
     10    1
     11    1
-    Name: Stage, dtype: int64
+    Name: Stage, dtype: int16
 
     >>> # Calculate the summary sleep statistics
     >>> hyp.sleep_statistics()
@@ -149,15 +151,15 @@ class Hypnogram:
     SLEEP          1      6
 
     All these methods and properties are also valid with a 5-stages hypnogram. In the example below,
-    we use the :py:func:`yasa.simulate_hypno` to generate a plausible 5-stages hypnogram with a
+    we use the :py:func:`yasa.simulate_hypnogram` to generate a plausible 5-stages hypnogram with a
     30-seconds resolution. A random seed is specified to ensure that we get reproducible results.
     Lastly, we set an actual start time to the hypnogram. As a result, the index of the resulting
     hypnogram is a :py:class:`pandas.DatetimeIndex`.
 
-    >>> from yasa import simulate_hypno
-    >>> hyp = simulate_hypno(
+    >>> from yasa import simulate_hypnogram
+    >>> hyp = simulate_hypnogram(
     ...     tib=500, n_stages=5, start="2022-12-15 22:30:00", scorer="S1", seed=42)
-    >>> hyp
+    >>> hyp  # This is a shortcut to `hyp.hypno`
     Time
     2022-12-15 22:30:00    WAKE
     2022-12-15 22:30:30    WAKE
@@ -170,7 +172,8 @@ class Hypnogram:
     2022-12-16 06:48:30      N2
     2022-12-16 06:49:00      N2
     2022-12-16 06:49:30      N2
-    Freq: 30S, Name: S1, Length: 1000, dtype: object
+    Freq: 30S, Name: S1, Length: 1000, dtype: category
+    Categories (7, object): ['WAKE', 'N1', 'N2', 'N3', 'REM', 'ART', 'UNS']
 
     The summary sleep statistics will include more items with a 5-stages hypnogram than a 2-stages
     hypnogram, i.e. the amount and percentage of each sleep stage, the REM latency, etc.
@@ -238,6 +241,10 @@ class Hypnogram:
         map_accepted = {"S": "SLEEP", "W": "WAKE", "R": "REM"}
         hypno = hypno.replace(map_accepted)
         labels = pd.Series(accepted).replace(map_accepted).unique().tolist()
+        # Change dtype of series to "categorical" (reduces memory)
+        cat_dtype = CategoricalDtype(labels, ordered=False)
+        hypno = hypno.astype(cat_dtype)
+        # Create Index
         if start is not None:
             hypno.index = pd.date_range(start=start, freq=freq, periods=hypno.shape[0])
             hypno.index.name = "Time"
@@ -246,6 +253,7 @@ class Hypnogram:
             fake_dt = pd.date_range(start="2022-12-03 00:00:00", freq=freq, periods=hypno.shape[0])
             hypno.index.name = "Epoch"
             timedelta = fake_dt - fake_dt[0]
+        # Set attributes
         self._hypno = hypno
         self._n_epochs = hypno.shape[0]
         self._freq = freq
@@ -266,7 +274,14 @@ class Hypnogram:
 
     @property
     def hypno(self):
-        """The hypnogram values, stored in a :py:class:`pandas.Series`."""
+        """
+        The hypnogram values, stored in a :py:class:`pandas.Series`. To reduce memory usage, the
+        stages are stored as categories (:py:class:`pandas.Categorical`). ``hypno``
+        inherits all the methods of a standard :py:class:`pandas.Series`, e.g. ``.describe()``,
+        ``.unique()``, ``.to_csv()``, and more.
+
+        .. note:: ``print(Hypnogram)`` is a shortcut to ``print(Hypnogram.hypno)``.
+        """
         return self._hypno
 
     @property
@@ -306,7 +321,8 @@ class Hypnogram:
     def n_stages(self):
         """
         The number of allowed stages in the hypnogram. This is not the number of unique stages
-        in the current hypnogram.
+        in the current hypnogram. This does not include Artefact and Unscored which are always
+        allowed.
         """
         return self._n_stages
 
@@ -345,16 +361,6 @@ class Hypnogram:
 
     # CLASS METHODS BELOW
 
-    def copy(self):
-        """Return a new copy of the current Hypnogram."""
-        return type(self)(
-            values=self.hypno.to_numpy(),
-            n_stages=self.n_stages,
-            freq=self.freq,
-            start=self.start,
-            scorer=self.scorer,
-        )
-
     def as_annotations(self):
         """
         Return a pandas DataFrame summarizing epoch-level information.
@@ -397,7 +403,7 @@ class Hypnogram:
         return pd.DataFrame(data).set_index("epoch")
 
     def as_int(self):
-        """Return hypnogram as integers.
+        """Return hypnogram values as integers.
 
         The default mapping from string to integer is:
 
@@ -412,7 +418,7 @@ class Hypnogram:
 
         Examples
         --------
-        Convert a 2-stages hypnogram to integers
+        Convert a 2-stages hypnogram to a pandas.Series of integers
 
         >>> from yasa import Hypnogram
         >>> hyp = Hypnogram(["W", "W", "S", "S", "W", "S"], n_stages=2)
@@ -424,7 +430,7 @@ class Hypnogram:
         3    1
         4    0
         5    1
-        Name: Stage, dtype: int64
+        Name: Stage, dtype: int16
 
         Same with a 4-stages hypnogram
 
@@ -439,9 +445,10 @@ class Hypnogram:
         4    3
         5    4
         6    0
-        Name: Stage, dtype: int64
+        Name: Stage, dtype: int16
         """
-        return self.hypno.replace(self.mapping).astype(int)
+        # Return as int16 (-32768 to 32767) to reduce memory usage
+        return self.hypno.replace(self.mapping).astype(np.int16)
 
     def consolidate_stages(self, new_n_stages):
         """Reduce the number of stages in a hypnogram to match actigraphy or wearables.
@@ -486,7 +493,8 @@ class Hypnogram:
         5    SLEEP
         6    SLEEP
         7     WAKE
-        Name: Stage, dtype: object
+        Name: Stage, dtype: category
+        Categories (4, object): ['WAKE', 'SLEEP', 'ART', 'UNS']
         """
         assert self.n_stages in [3, 4, 5], "`self.n_stages` must be 3, 4, or 5"
         assert new_n_stages in [2, 3, 4], "`new_n_stages` must be 2, 3, or 4"
@@ -520,6 +528,16 @@ class Hypnogram:
             scorer=self.scorer,
         )
 
+    def copy(self):
+        """Return a new copy of the current Hypnogram."""
+        return type(self)(
+            values=self.hypno.to_numpy(),
+            n_stages=self.n_stages,
+            freq=self.freq,
+            start=self.start,
+            scorer=self.scorer,
+        )
+
     def find_periods(self, threshold="5min", equal_length=False):
         """Find sequences of consecutive values exceeding a certain duration in hypnogram.
 
@@ -579,8 +597,8 @@ class Hypnogram:
         This function is not limited to binary arrays, e.g. a 5-stages hypnogram at 30-sec
         resolution:
 
-        >>> from yasa import simulate_hypno
-        >>> hyp = simulate_hypno(tib=30, seed=42)
+        >>> from yasa import simulate_hypnogram
+        >>> hyp = simulate_hypnogram(tib=30, seed=42)
         >>> hyp.find_periods(threshold="2min")
           values  start  length
         0   WAKE      0       5
@@ -626,21 +644,21 @@ class Hypnogram:
         .. plot::
 
             >>> from yasa import simulate_hypnogram
-            >>> ax = simulate_hypno(tib=480, seed=88).plot_hypnogram(highlight="REM"")
+            >>> ax = simulate_hypnogram(tib=480, seed=88).plot_hypnogram(highlight="REM")
         """
         return plot_hypnogram(self, **kwargs)
 
     def simulate_similar(self, **kwargs):
         """Simulate a new hypnogram based on properties of the current hypnogram.
 
-        .. seealso:: :py:func:`yasa.simulate_hypno`
+        .. seealso:: :py:func:`yasa.simulate_hypnogram`
 
         Parameters
         ----------
         self : :py:class:`yasa.Hypnogram`
             Hypnogram, assumed to be already cropped to time in bed (TIB).
         **kwargs : dict
-            Optional keyword arguments passed to :py:func:`yasa.simulate_hypno`.
+            Optional keyword arguments passed to :py:func:`yasa.simulate_hypnogram`.
 
         Returns
         -------
@@ -675,7 +693,7 @@ class Hypnogram:
             "`n_stages` and `freq` cannot be included as additional `**kwargs` "
             "because they must match properties of the current Hypnogram."
         )
-        simulate_hypno_kwargs = {
+        simulate_hypnogram_kwargs = {
             "tib": self.duration,
             "n_stages": self.n_stages,
             "freq": self.freq,
@@ -683,8 +701,8 @@ class Hypnogram:
             "start": self.start,
             "scorer": self.scorer,
         }
-        simulate_hypno_kwargs.update(kwargs)
-        return simulate_hypno(**simulate_hypno_kwargs)
+        simulate_hypnogram_kwargs.update(kwargs)
+        return simulate_hypnogram(**simulate_hypnogram_kwargs)
 
     def sleep_statistics(self):
         """
@@ -755,9 +773,9 @@ class Hypnogram:
 
         Sleep statistics for a 5-stages hypnogram
 
-        >>> from yasa import simulate_hypno
+        >>> from yasa import simulate_hypnogram
         >>> # Generate a 8 hr (= 480 minutes) 5-stages hypnogram with a 30-seconds resolution
-        >>> hyp = simulate_hypno(tib=480, seed=42)
+        >>> hyp = simulate_hypnogram(tib=480, seed=42)
         >>> hyp.sleep_statistics()
         {'TIB': 480.0,
         'SPT': 477.5,
@@ -886,9 +904,9 @@ class Hypnogram:
 
         Examples
         --------
-        >>> from yasa import Hypnogram, simulate_hypno
+        >>> from yasa import Hypnogram, simulate_hypnogram
         >>> # Generate a 8 hr (= 480 minutes) 5-stages hypnogram with a 30-seconds resolution
-        >>> hyp = simulate_hypno(tib=480, seed=42)
+        >>> hyp = simulate_hypnogram(tib=480, seed=42)
         >>> counts, probs = hyp.transition_matrix()
         >>> counts
         To Stage    WAKE  N1   N2  N3  REM
@@ -946,7 +964,8 @@ class Hypnogram:
         2022-12-23 23:01:00    SLEEP
         2022-12-23 23:01:30    SLEEP
         2022-12-23 23:02:00     WAKE
-        Freq: 30S, Name: Stage, dtype: object
+        Freq: 30S, Name: Stage, dtype: category
+        Categories (4, object): ['WAKE', 'SLEEP', 'ART', 'UNS']
 
         Upsample to a 15-seconds resolution
 
@@ -963,7 +982,8 @@ class Hypnogram:
         2022-12-23 23:01:45    SLEEP
         2022-12-23 23:02:00     WAKE
         2022-12-23 23:02:15     WAKE
-        Freq: 15S, Name: Stage, dtype: object
+        Freq: 15S, Name: Stage, dtype: category
+        Categories (4, object): ['WAKE', 'SLEEP', 'ART', 'UNS']
         """
         assert pd.Timedelta(new_freq) < pd.Timedelta(self.freq), (
             f"The upsampling `new_freq` ({new_freq}) must be higher than the current frequency of "
@@ -1497,7 +1517,7 @@ def hypno_find_periods(hypno, sf_hypno, threshold="5min", equal_length=False):
 #############################################################################
 
 
-def simulate_hypno(
+def simulate_hypnogram(
     tib=480,
     trans_probas=None,
     init_probas=None,
@@ -1577,8 +1597,8 @@ def simulate_hypno(
 
     Examples
     --------
-    >>> from yasa import simulate_hypno
-    >>> hyp = simulate_hypno(tib=5, seed=1)
+    >>> from yasa import simulate_hypnogram
+    >>> hyp = simulate_hypnogram(tib=5, seed=1)
     >>> print(hyp)
     Epoch
     0    WAKE
@@ -1593,7 +1613,7 @@ def simulate_hypno(
     9      N2
     Name: Stage, dtype: object
 
-    >>> hyp = simulate_hypno(tib=5, n_stages=2, seed=1)
+    >>> hyp = simulate_hypnogram(tib=5, n_stages=2, seed=1)
     >>> print(hyp)
     Epoch
     0     WAKE
@@ -1610,7 +1630,7 @@ def simulate_hypno(
 
     Add some Unscored epochs.
 
-    >>> hyp = simulate_hypno(tib=5, n_stages=2, seed=1)
+    >>> hyp = simulate_hypnogram(tib=5, n_stages=2, seed=1)
     >>> hyp.hypno.iloc[-2:] = "UNS"
     >>> print(hyp)
     Epoch
diff --git a/yasa/plotting.py b/yasa/plotting.py
index 39d08c7..87ba2f0 100644
--- a/yasa/plotting.py
+++ b/yasa/plotting.py
@@ -41,9 +41,9 @@ def plot_hypnogram(hyp, lw=1.5, highlight="REM", fill_color=None, ax=None):
     --------
     .. plot::
 
-        >>> from yasa import simulate_hypno
+        >>> from yasa import simulate_hypnogram
         >>> import matplotlib.pyplot as plt
-        >>> hyp = simulate_hypno(tib=300, seed=11)
+        >>> hyp = simulate_hypnogram(tib=300, seed=11)
         >>> ax = hyp.plot_hypnogram()
         >>> plt.tight_layout()
 
@@ -58,8 +58,8 @@ def plot_hypnogram(hyp, lw=1.5, highlight="REM", fill_color=None, ax=None):
     .. plot::
 
         >>> fig, axes = plt.subplots(nrows=2, figsize=(6, 4), constrained_layout=True)
-        >>> hyp_a = simulate_hypno(n_stages=3, seed=99)
-        >>> hyp_b = simulate_hypno(n_stages=3, seed=99, start="2022-01-31 23:30:00")
+        >>> hyp_a = simulate_hypnogram(n_stages=3, seed=99)
+        >>> hyp_b = simulate_hypnogram(n_stages=3, seed=99, start="2022-01-31 23:30:00")
         >>> hyp_a.plot_hypnogram(lw=1, fill_color="whitesmoke", highlight=None, ax=axes[0])
         >>> hyp_b.plot_hypnogram(lw=1, fill_color="whitesmoke", highlight=None, ax=axes[1])
     """


Test patch:
diff --git a/yasa/tests/test_hypno.py b/yasa/tests/test_hypno.py
index 132e1db..1c8a2e4 100644
--- a/yasa/tests/test_hypno.py
+++ b/yasa/tests/test_hypno.py
@@ -11,7 +11,7 @@ from yasa.hypno import (
     hypno_upsample_to_sf,
     hypno_fit_to_data,
     hypno_upsample_to_data,
-    simulate_hypno,
+    simulate_hypnogram,
 )
 
 from yasa.hypno import hypno_find_periods as hfp
@@ -131,26 +131,28 @@ class TestHypno(unittest.TestCase):
 
     def test_simulation(self):
         """Test hypnogram simulations."""
-        hyp = simulate_hypno(tib=4, seed=1)
+        hyp = simulate_hypnogram(tib=4, seed=1)
         assert hyp.n_epochs == 8
         np.testing.assert_array_equal(hyp.as_int(), [0, 1, 1, 2, 2, 2, 2, 2])
 
         # Handling different n_stages
-        assert simulate_hypno(tib=1000, n_stages=2).hypno.nunique() == 2
+        assert simulate_hypnogram(tib=1000, n_stages=2).hypno.nunique() == 2
         np.testing.assert_array_equal(
-            simulate_hypno(tib=4, seed=1, n_stages=3).as_int(), [0, 2, 2, 2, 2, 2, 2, 2]
+            simulate_hypnogram(tib=4, seed=1, n_stages=3).as_int(), [0, 2, 2, 2, 2, 2, 2, 2]
         )
 
         # Handling different frequencies
-        np.testing.assert_array_equal(hyp.hypno, simulate_hypno(tib=4, freq="0.5min", seed=1).hypno)
         np.testing.assert_array_equal(
-            hyp.upsample("5s").hypno, simulate_hypno(tib=4, freq="5s", seed=1).hypno
+            hyp.hypno, simulate_hypnogram(tib=4, freq="0.5min", seed=1).hypno
         )
-        assert simulate_hypno(tib=4, freq="15s").n_epochs == 16
-        assert simulate_hypno(tib=4, freq="15s").duration == 4
-        assert simulate_hypno(tib=4, freq="30s").duration == 4
+        np.testing.assert_array_equal(
+            hyp.upsample("5s").hypno, simulate_hypnogram(tib=4, freq="5s", seed=1).hypno
+        )
+        assert simulate_hypnogram(tib=4, freq="15s").n_epochs == 16
+        assert simulate_hypnogram(tib=4, freq="15s").duration == 4
+        assert simulate_hypnogram(tib=4, freq="30s").duration == 4
         with pytest.raises(AssertionError):
-            simulate_hypno(freq="60s")
+            simulate_hypnogram(freq="60s")
 
         # Hanling different probabilities
         trans_probas = pd.DataFrame(
@@ -158,22 +160,22 @@ class TestHypno(unittest.TestCase):
             index=["WAKE", "N1", "N2", "N3", "REM"],
             columns=["WAKE", "N1", "N2", "N3", "REM"],
         )
-        simulate_hypno(tib=2, trans_probas=trans_probas)
-        simulate_hypno(tib=2, init_probas=trans_probas.loc["WAKE"])
-        simulate_hypno(tib=2, trans_probas=trans_probas, init_probas=trans_probas.loc["WAKE"])
+        simulate_hypnogram(tib=2, trans_probas=trans_probas)
+        simulate_hypnogram(tib=2, init_probas=trans_probas.loc["WAKE"])
+        simulate_hypnogram(tib=2, trans_probas=trans_probas, init_probas=trans_probas.loc["WAKE"])
         # Setting all probabilities between stages as zero
         trans_probas.loc[:, :] = np.eye(5, 5)
-        assert not simulate_hypno(trans_probas=trans_probas).as_int().any()
+        assert not simulate_hypnogram(trans_probas=trans_probas).as_int().any()
 
         # When trans_proba has more stages than allowed by n_stages
         with pytest.raises(AssertionError):
-            simulate_hypno(n_stages=4, trans_probas=trans_probas)
+            simulate_hypnogram(n_stages=4, trans_probas=trans_probas)
 
         # When trans_probas includes only a subset of stages allowed by n_stages
         trans_probas = trans_probas.drop(index="REM", columns="REM")
         trans_probas.loc[:, :] = np.full((4, 4), 0.25)
-        simulate_hypno(trans_probas=trans_probas)
+        simulate_hypnogram(trans_probas=trans_probas)
 
         # Passing **kwargs through to yasa.Hypnogram
-        shyp = simulate_hypno(tib=5, scorer="RV", start="2022-12-15 22:30:00")
+        shyp = simulate_hypnogram(tib=5, scorer="RV", start="2022-12-15 22:30:00")
         assert shyp.scorer == shyp.hypno.name == "RV"
diff --git a/yasa/tests/test_hypnoclass.py b/yasa/tests/test_hypnoclass.py
index 7231981..74d166d 100644
--- a/yasa/tests/test_hypnoclass.py
+++ b/yasa/tests/test_hypnoclass.py
@@ -1,10 +1,11 @@
 """Test the class Hypnogram."""
 import mne
+import pytest
 import unittest
 import numpy as np
 import pandas as pd
 import matplotlib.pyplot as plt
-from yasa.hypno import simulate_hypno, Hypnogram, hypno_str_to_int
+from yasa.hypno import simulate_hypnogram, Hypnogram, hypno_str_to_int
 
 
 def create_raw(npts, ch_names=["F4-M1", "F3-M2"], sf=100):
@@ -21,13 +22,14 @@ class TestHypnoClass(unittest.TestCase):
 
     def test_2stages_hypno(self):
         """Test 2-stages Hypnogram class"""
-        hyp = simulate_hypno(tib=120, n_stages=2, seed=42)
+        hyp = simulate_hypnogram(tib=120, n_stages=2, seed=42)
         print(hyp)
         print(str(hyp))
 
         # Check properties
         np.testing.assert_array_equal(hyp.hypno.str.get(0)[:10], np.repeat(["W", "S"], 5))
         assert isinstance(hyp.hypno.index, pd.RangeIndex)
+        assert hyp.hypno.dtype == "category"
         assert hyp.hypno.index.name == "Epoch"
         assert hyp.sampling_frequency == 1 / 30
         assert hyp.freq == "30s"
@@ -104,7 +106,7 @@ class TestHypnoClass(unittest.TestCase):
         hyp_up = hyp.upsample_to_data(raw)
         assert isinstance(hyp_up, np.ndarray)
         assert hyp_up.size == npts
-        assert hyp_up.dtype == int
+        assert hyp_up.dtype == np.int16
         hyp_up = hyp.upsample_to_data(raw.get_data(), sf=100)
 
         # yasa.Hypnogram.simulate_similar
@@ -121,7 +123,7 @@ class TestHypnoClass(unittest.TestCase):
         assert hyp.simulate_similar(tib=2, scorer="YASA").scorer == "YASA"
         assert hyp.simulate_similar(tib=2, start="2022-11-10").start == "2022-11-10"
         np.testing.assert_array_equal(
-            simulate_hypno(seed=1).simulate_similar(tib=5, seed=6).as_int(),
+            simulate_hypnogram(seed=1).simulate_similar(tib=5, seed=6).as_int(),
             [0, 0, 0, 0, 1, 1, 1, 2, 2, 2],
         )
 
@@ -134,7 +136,7 @@ class TestHypnoClass(unittest.TestCase):
 
     def test_3stages_hypno(self):
         """Test 3-stages Hypnogram class"""
-        hyp = simulate_hypno(tib=120, n_stages=3, freq="1s", seed=42)
+        hyp = simulate_hypnogram(tib=120, n_stages=3, freq="1s", seed=42)
         assert hyp.sampling_frequency == 1
         assert hyp.freq == "1s"
         assert hyp.n_stages == 3
@@ -145,9 +147,13 @@ class TestHypnoClass(unittest.TestCase):
         assert "%REM" in sstats.keys()
         assert "Lat_REM" in sstats.keys()
 
+        # Try to set a value that is not a valid category
+        with pytest.raises((TypeError, ValueError)):  # TypeError in newer versions of Pandas
+            hyp.hypno.loc[0] = "Dream sleep"
+
     def test_4stages_hypno(self):
         """Test 4-stages Hypnogram class"""
-        hyp = simulate_hypno(tib=400, n_stages=4, freq="30s", seed=42)
+        hyp = simulate_hypnogram(tib=400, n_stages=4, freq="30s", seed=42)
         assert hyp.n_stages == 4
         assert hyp.labels == ["WAKE", "LIGHT", "DEEP", "REM", "ART", "UNS"]
         assert hyp.mapping == {"WAKE": 0, "LIGHT": 2, "DEEP": 3, "REM": 4, "ART": -1, "UNS": -2}
@@ -159,7 +165,7 @@ class TestHypnoClass(unittest.TestCase):
 
     def test_5stages_hypno(self):
         """Test 5-stages Hypnogram class"""
-        hyp = simulate_hypno(tib=600, n_stages=5, freq="30s", seed=42)
+        hyp = simulate_hypnogram(tib=600, n_stages=5, freq="30s", seed=42)
         assert hyp.n_stages == 5
         assert hyp.labels == ["WAKE", "N1", "N2", "N3", "REM", "ART", "UNS"]
         assert hyp.mapping == {
diff --git a/yasa/tests/test_plotting.py b/yasa/tests/test_plotting.py
index b8ae5e2..cb38b1b 100644
--- a/yasa/tests/test_plotting.py
+++ b/yasa/tests/test_plotting.py
@@ -4,7 +4,7 @@ import unittest
 import numpy as np
 import pandas as pd
 import matplotlib.pyplot as plt
-from yasa.hypno import simulate_hypno
+from yasa.hypno import simulate_hypnogram
 from yasa.plotting import topoplot, plot_hypnogram
 
 
@@ -41,8 +41,8 @@ class TestPlotting(unittest.TestCase):
         with pytest.raises(AssertionError):
             _ = plot_hypnogram(np.repeat([0, 1, 2, 3, 4, -2, -1, -3], 120))
         # Default parameters
-        hyp5 = simulate_hypno(n_stages=5)
-        hyp2 = simulate_hypno(n_stages=2)
+        hyp5 = simulate_hypnogram(n_stages=5)
+        hyp2 = simulate_hypnogram(n_stages=2)
         ax = hyp5.plot_hypnogram()
         assert isinstance(ax, plt.Axes)
         # Aesthetic parameters
@@ -54,7 +54,7 @@ class TestPlotting(unittest.TestCase):
         ax = plt.subplot()
         _ = hyp5.plot_hypnogram(ax=ax)
         # With datetime axis
-        hyp3 = simulate_hypno(n_stages=3, tib=800, start="2020-01-01 20:00:00")
+        hyp3 = simulate_hypnogram(n_stages=3, tib=800, start="2020-01-01 20:00:00")
         hyp3.plot_hypnogram()
         # With Artefacts and Unscored
         hyp3.hypno.iloc[-100:] = "UNS"