diff --git a/autokoopman/_version.py b/autokoopman/_version.py
index 30f93da..df51d96 100644
--- a/autokoopman/_version.py
+++ b/autokoopman/_version.py
@@ -1,4 +1,3 @@
-
 # This file helps to compute a version number in source trees obtained from
 # git-archive tarball (such as those provided by githubs download-from-tag
 # feature). Distribution tarballs (built by setup.py sdist) and build
@@ -61,17 +60,18 @@ class NotThisMethod(Exception):
 
 def register_vcs_handler(vcs, method):  # decorator
     """Create decorator to mark a method as the handler of a VCS."""
+
     def decorate(f):
         """Store f in HANDLERS[vcs][method]."""
         if vcs not in HANDLERS:
             HANDLERS[vcs] = {}
         HANDLERS[vcs][method] = f
         return f
+
     return decorate
 
 
-def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False,
-                env=None):
+def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, env=None):
     """Call the given command(s)."""
     assert isinstance(commands, list)
     process = None
@@ -87,10 +87,14 @@ def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False,
         try:
             dispcmd = str([command] + args)
             # remember shell=False, so use git.cmd on windows, not just git
-            process = subprocess.Popen([command] + args, cwd=cwd, env=env,
-                                       stdout=subprocess.PIPE,
-                                       stderr=(subprocess.PIPE if hide_stderr
-                                               else None), **popen_kwargs)
+            process = subprocess.Popen(
+                [command] + args,
+                cwd=cwd,
+                env=env,
+                stdout=subprocess.PIPE,
+                stderr=(subprocess.PIPE if hide_stderr else None),
+                **popen_kwargs,
+            )
             break
         except OSError:
             e = sys.exc_info()[1]
@@ -125,15 +129,21 @@ def versions_from_parentdir(parentdir_prefix, root, verbose):
     for _ in range(3):
         dirname = os.path.basename(root)
         if dirname.startswith(parentdir_prefix):
-            return {"version": dirname[len(parentdir_prefix):],
-                    "full-revisionid": None,
-                    "dirty": False, "error": None, "date": None}
+            return {
+                "version": dirname[len(parentdir_prefix) :],
+                "full-revisionid": None,
+                "dirty": False,
+                "error": None,
+                "date": None,
+            }
         rootdirs.append(root)
         root = os.path.dirname(root)  # up a level
 
     if verbose:
-        print("Tried directories %s but none started with prefix %s" %
-              (str(rootdirs), parentdir_prefix))
+        print(
+            "Tried directories %s but none started with prefix %s"
+            % (str(rootdirs), parentdir_prefix)
+        )
     raise NotThisMethod("rootdir doesn't start with parentdir_prefix")
 
 
@@ -192,7 +202,7 @@ def git_versions_from_keywords(keywords, tag_prefix, verbose):
     # starting in git-1.8.3, tags are listed as "tag: foo-1.0" instead of
     # just "foo-1.0". If we see a "tag: " prefix, prefer those.
     TAG = "tag: "
-    tags = {r[len(TAG):] for r in refs if r.startswith(TAG)}
+    tags = {r[len(TAG) :] for r in refs if r.startswith(TAG)}
     if not tags:
         # Either we're using git < 1.8.3, or there really are no tags. We use
         # a heuristic: assume all version tags have a digit. The old git %d
@@ -201,7 +211,7 @@ def git_versions_from_keywords(keywords, tag_prefix, verbose):
         # between branches and tags. By ignoring refnames without digits, we
         # filter out many common branch names like "release" and
         # "stabilization", as well as "HEAD" and "master".
-        tags = {r for r in refs if re.search(r'\d', r)}
+        tags = {r for r in refs if re.search(r"\d", r)}
         if verbose:
             print("discarding '%s', no digits" % ",".join(refs - tags))
     if verbose:
@@ -209,24 +219,31 @@ def git_versions_from_keywords(keywords, tag_prefix, verbose):
     for ref in sorted(tags):
         # sorting will prefer e.g. "2.0" over "2.0rc1"
         if ref.startswith(tag_prefix):
-            r = ref[len(tag_prefix):]
+            r = ref[len(tag_prefix) :]
             # Filter out refs that exactly match prefix or that don't start
             # with a number once the prefix is stripped (mostly a concern
             # when prefix is '')
-            if not re.match(r'\d', r):
+            if not re.match(r"\d", r):
                 continue
             if verbose:
                 print("picking %s" % r)
-            return {"version": r,
-                    "full-revisionid": keywords["full"].strip(),
-                    "dirty": False, "error": None,
-                    "date": date}
+            return {
+                "version": r,
+                "full-revisionid": keywords["full"].strip(),
+                "dirty": False,
+                "error": None,
+                "date": date,
+            }
     # no suitable tags, so version is "0+unknown", but full hex is still there
     if verbose:
         print("no suitable tags, using unknown + full revision id")
-    return {"version": "0+unknown",
-            "full-revisionid": keywords["full"].strip(),
-            "dirty": False, "error": "no suitable tags", "date": None}
+    return {
+        "version": "0+unknown",
+        "full-revisionid": keywords["full"].strip(),
+        "dirty": False,
+        "error": "no suitable tags",
+        "date": None,
+    }
 
 
 @register_vcs_handler("git", "pieces_from_vcs")
@@ -248,8 +265,7 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command):
     env.pop("GIT_DIR", None)
     runner = functools.partial(runner, env=env)
 
-    _, rc = runner(GITS, ["rev-parse", "--git-dir"], cwd=root,
-                   hide_stderr=not verbose)
+    _, rc = runner(GITS, ["rev-parse", "--git-dir"], cwd=root, hide_stderr=not verbose)
     if rc != 0:
         if verbose:
             print("Directory %s not under git control" % root)
@@ -257,10 +273,19 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command):
 
     # if there is a tag matching tag_prefix, this yields TAG-NUM-gHEX[-dirty]
     # if there isn't one, this yields HEX[-dirty] (no NUM)
-    describe_out, rc = runner(GITS, [
-        "describe", "--tags", "--dirty", "--always", "--long",
-        "--match", f"{tag_prefix}[[:digit:]]*"
-    ], cwd=root)
+    describe_out, rc = runner(
+        GITS,
+        [
+            "describe",
+            "--tags",
+            "--dirty",
+            "--always",
+            "--long",
+            "--match",
+            f"{tag_prefix}[[:digit:]]*",
+        ],
+        cwd=root,
+    )
     # --long was added in git-1.5.5
     if describe_out is None:
         raise NotThisMethod("'git describe' failed")
@@ -275,8 +300,7 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command):
     pieces["short"] = full_out[:7]  # maybe improved later
     pieces["error"] = None
 
-    branch_name, rc = runner(GITS, ["rev-parse", "--abbrev-ref", "HEAD"],
-                             cwd=root)
+    branch_name, rc = runner(GITS, ["rev-parse", "--abbrev-ref", "HEAD"], cwd=root)
     # --abbrev-ref was added in git-1.6.3
     if rc != 0 or branch_name is None:
         raise NotThisMethod("'git rev-parse --abbrev-ref' returned error")
@@ -316,17 +340,16 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command):
     dirty = git_describe.endswith("-dirty")
     pieces["dirty"] = dirty
     if dirty:
-        git_describe = git_describe[:git_describe.rindex("-dirty")]
+        git_describe = git_describe[: git_describe.rindex("-dirty")]
 
     # now we have TAG-NUM-gHEX or HEX
 
     if "-" in git_describe:
         # TAG-NUM-gHEX
-        mo = re.search(r'^(.+)-(\d+)-g([0-9a-f]+)$', git_describe)
+        mo = re.search(r"^(.+)-(\d+)-g([0-9a-f]+)$", git_describe)
         if not mo:
             # unparsable. Maybe git-describe is misbehaving?
-            pieces["error"] = ("unable to parse git-describe output: '%s'"
-                               % describe_out)
+            pieces["error"] = "unable to parse git-describe output: '%s'" % describe_out
             return pieces
 
         # tag
@@ -335,10 +358,12 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command):
             if verbose:
                 fmt = "tag '%s' doesn't start with prefix '%s'"
                 print(fmt % (full_tag, tag_prefix))
-            pieces["error"] = ("tag '%s' doesn't start with prefix '%s'"
-                               % (full_tag, tag_prefix))
+            pieces["error"] = "tag '%s' doesn't start with prefix '%s'" % (
+                full_tag,
+                tag_prefix,
+            )
             return pieces
-        pieces["closest-tag"] = full_tag[len(tag_prefix):]
+        pieces["closest-tag"] = full_tag[len(tag_prefix) :]
 
         # distance: number of commits since tag
         pieces["distance"] = int(mo.group(2))
@@ -387,8 +412,7 @@ def render_pep440(pieces):
                 rendered += ".dirty"
     else:
         # exception #1
-        rendered = "0+untagged.%d.g%s" % (pieces["distance"],
-                                          pieces["short"])
+        rendered = "0+untagged.%d.g%s" % (pieces["distance"], pieces["short"])
         if pieces["dirty"]:
             rendered += ".dirty"
     return rendered
@@ -417,8 +441,7 @@ def render_pep440_branch(pieces):
         rendered = "0"
         if pieces["branch"] != "master":
             rendered += ".dev0"
-        rendered += "+untagged.%d.g%s" % (pieces["distance"],
-                                          pieces["short"])
+        rendered += "+untagged.%d.g%s" % (pieces["distance"], pieces["short"])
         if pieces["dirty"]:
             rendered += ".dirty"
     return rendered
@@ -579,11 +602,13 @@ def render_git_describe_long(pieces):
 def render(pieces, style):
     """Render the given version pieces into the requested style."""
     if pieces["error"]:
-        return {"version": "unknown",
-                "full-revisionid": pieces.get("long"),
-                "dirty": None,
-                "error": pieces["error"],
-                "date": None}
+        return {
+            "version": "unknown",
+            "full-revisionid": pieces.get("long"),
+            "dirty": None,
+            "error": pieces["error"],
+            "date": None,
+        }
 
     if not style or style == "default":
         style = "pep440"  # the default
@@ -607,9 +632,13 @@ def render(pieces, style):
     else:
         raise ValueError("unknown style '%s'" % style)
 
-    return {"version": rendered, "full-revisionid": pieces["long"],
-            "dirty": pieces["dirty"], "error": None,
-            "date": pieces.get("date")}
+    return {
+        "version": rendered,
+        "full-revisionid": pieces["long"],
+        "dirty": pieces["dirty"],
+        "error": None,
+        "date": pieces.get("date"),
+    }
 
 
 def get_versions():
@@ -623,8 +652,7 @@ def get_versions():
     verbose = cfg.verbose
 
     try:
-        return git_versions_from_keywords(get_keywords(), cfg.tag_prefix,
-                                          verbose)
+        return git_versions_from_keywords(get_keywords(), cfg.tag_prefix, verbose)
     except NotThisMethod:
         pass
 
@@ -633,13 +661,16 @@ def get_versions():
         # versionfile_source is the relative path from the top of the source
         # tree (where the .git directory might live) to this file. Invert
         # this to find the root from __file__.
-        for _ in cfg.versionfile_source.split('/'):
+        for _ in cfg.versionfile_source.split("/"):
             root = os.path.dirname(root)
     except NameError:
-        return {"version": "0+unknown", "full-revisionid": None,
-                "dirty": None,
-                "error": "unable to find root of source tree",
-                "date": None}
+        return {
+            "version": "0+unknown",
+            "full-revisionid": None,
+            "dirty": None,
+            "error": "unable to find root of source tree",
+            "date": None,
+        }
 
     try:
         pieces = git_pieces_from_vcs(cfg.tag_prefix, root, verbose)
@@ -653,6 +684,10 @@ def get_versions():
     except NotThisMethod:
         pass
 
-    return {"version": "0+unknown", "full-revisionid": None,
-            "dirty": None,
-            "error": "unable to compute version", "date": None}
+    return {
+        "version": "0+unknown",
+        "full-revisionid": None,
+        "dirty": None,
+        "error": "unable to compute version",
+        "date": None,
+    }
diff --git a/autokoopman/autokoopman.py b/autokoopman/autokoopman.py
index f247f0f..9a210cb 100644
--- a/autokoopman/autokoopman.py
+++ b/autokoopman/autokoopman.py
@@ -103,15 +103,24 @@ def get_parameter_space(obs_type, threshold_range, rank):
 
 
 def get_estimator(
-    obs_type, learn_continuous, sampling_period, dim, obs, hyperparams, normalize
+    obs_type,
+    learn_continuous,
+    sampling_period,
+    dim,
+    obs,
+    hyperparams,
+    normalize,
+    weights,
 ):
     """from the myriad of user suppled switches, select the right estimator"""
 
     def inst_estimator(*args, **kwargs):
         if learn_continuous:
-            return KoopmanContinuousEstimator(args[0], *args[2:], **kwargs)
+            return KoopmanContinuousEstimator(
+                args[0], *args[2:], weights=weights, **kwargs
+            )
         else:
-            return KoopmanDiscEstimator(*args, **kwargs)
+            return KoopmanDiscEstimator(*args, weights=weights, **kwargs)
 
     if obs_type == "rff":
         observables = kobs.IdentityObservable() | kobs.RFFObservable(
@@ -160,6 +169,9 @@ def auto_koopman(
     scoring_weights: Optional[
         Union[Sequence[np.ndarray], Dict[Hashable, np.ndarray]]
     ] = None,
+    learning_weights: Optional[
+        Union[Sequence[np.ndarray], Dict[Hashable, np.ndarray]]
+    ] = None,
     n_obs: int = 100,
     rank: Optional[Union[Tuple[int, int], Tuple[int, int, int]]] = None,
     grid_param_slices: int = 10,
@@ -234,13 +246,19 @@ def auto_koopman(
             scoring_weights is not None
         ), f"weighted scoring requires scoring weights (currently None)"
 
-    training_data, sampling_period, scoring_weights = _sanitize_training_data(
+    (
+        training_data,
+        sampling_period,
+        scoring_weights,
+        learning_weights,
+    ) = _sanitize_training_data(
         training_data,
         inputs_training_data,
         sampling_period,
         opt,
         obs_type,
         scoring_weights,
+        learning_weights,
     )
 
     # get the hyperparameter map
@@ -267,6 +285,7 @@ def auto_koopman(
             lengthscale,
             sampling_period,
             normalize,
+            learning_weights,
         )
 
     # setup the tuner
@@ -376,6 +395,7 @@ def _edmd_model_map(
     lengthscale,
     sampling_period,
     normalize,
+    weights,
 ) -> HyperparameterMap:
     """model map for eDMD based methods
 
@@ -419,6 +439,7 @@ def get_model(self, hyperparams: Sequence):
                 n_obs,
                 hyperparams,
                 normalize,
+                weights,
             )
 
     # get the hyperparameter map
@@ -426,7 +447,13 @@ def get_model(self, hyperparams: Sequence):
 
 
 def _sanitize_training_data(
-    training_data, inputs_training_data, sampling_period, opt, obs_type, scoring_weights
+    training_data,
+    inputs_training_data,
+    sampling_period,
+    opt,
+    obs_type,
+    scoring_weights,
+    learning_weights,
 ):
     """auto_koopman input sanitization"""
 
@@ -450,27 +477,48 @@ def _sanitize_training_data(
     # convert the data to autokoopman trajectories
     if isinstance(training_data, TrajectoriesData):
         if not isinstance(training_data, UniformTimeTrajectoriesData):
-            assert scoring_weights is None, f"scoring weights must be None as interpolation is occuring"
+            assert (
+                scoring_weights is None and learning_weights is None
+            ), f"weights must be None as interpolation is occuring"
             print(
                 f"resampling trajectories as they need to be uniform time (sampling period {sampling_period})"
             )
             training_data = training_data.interp_uniform_time(sampling_period)
         else:
             if not np.isclose(training_data.sampling_period, sampling_period):
-                assert scoring_weights is None, f"scoring weights must be None as interpolation is occuring"
+                assert (
+                    scoring_weights is None and learning_weights is None
+                ), f"weights must be None as interpolation is occuring"
                 print(
                     f"resampling trajectories because the sampling periods differ (original {training_data.sampling_period}, new {sampling_period})"
                 )
                 training_data = training_data.interp_uniform_time(sampling_period)
     else:
-        if isinstance(training_data, dict):
-            assert isinstance(scoring_weights, dict), "training data has unordered keys, so scoring weights must be a dictionary with matching keys"
+        if isinstance(training_data, dict) and scoring_weights is not None:
+            assert isinstance(
+                scoring_weights, dict
+            ), "training data has unordered keys, so scoring weights must be a dictionary with matching keys"
         else:
-            scoring_weights = {idx: weights for idx, weights in enumerate(scoring_weights)}
+            if scoring_weights is not None and not isinstance(scoring_weights, dict):
+                scoring_weights = {
+                    idx: weights for idx, weights in enumerate(scoring_weights)
+                }
+
+        if isinstance(training_data, dict) and learning_weights is not None:
+            assert isinstance(
+                learning_weights, dict
+            ), "training data has unordered keys, so learning weights must be a dictionary with matching keys"
+        else:
+            if learning_weights is not None and not isinstance(learning_weights, dict):
+                learning_weights = {
+                    idx: weights for idx, weights in enumerate(learning_weights)
+                }
 
         # figure out how to add inputs
         training_iter = (
-            training_data.items() if isinstance(training_data, dict) else enumerate(training_data)
+            training_data.items()
+            if isinstance(training_data, dict)
+            else enumerate(training_data)
         )
         if inputs_training_data is not None:
             training_iter = [(n, x, inputs_training_data[n]) for n, x in training_iter]
@@ -491,4 +539,4 @@ def _sanitize_training_data(
                 }
             )
 
-    return training_data, sampling_period, scoring_weights
+    return training_data, sampling_period, scoring_weights, learning_weights
diff --git a/autokoopman/core/estimator.py b/autokoopman/core/estimator.py
index 703793e..acd4a8f 100644
--- a/autokoopman/core/estimator.py
+++ b/autokoopman/core/estimator.py
@@ -21,6 +21,7 @@ class TrajectoryEstimator(abc.ABC):
     :param normalize: apply MinMax normalization to the fit data
     :param feature_range: range for MinMax scaler
     """
+
     def __init__(self, normalize=True, feature_range=(-1, 1)) -> None:
         self.normalize = normalize
         if self.normalize:
@@ -47,13 +48,18 @@ class NextStepEstimator(TrajectoryEstimator):
     :param feature_range: range for MinMax scaler
     """
 
-    def __init__(self, normalize=True, feature_range=(-1, 1)) -> None:
+    def __init__(self, normalize=True, feature_range=(-1, 1), weights=None) -> None:
         super().__init__(normalize=normalize, feature_range=feature_range)
         self.names = None
+        self.weights = weights
 
     @abc.abstractmethod
     def fit_next_step(
-        self, X: np.ndarray, Y: np.ndarray, U: Optional[np.ndarray] = None
+        self,
+        X: np.ndarray,
+        Y: np.ndarray,
+        U: Optional[np.ndarray] = None,
+        weights: Optional[np.ndarray] = None,
     ) -> None:
         """an alternative fit method that uses a trajectories data structure"""
         pass
@@ -62,11 +68,15 @@ def fit(self, X: atraj.TrajectoriesData) -> None:
         assert isinstance(
             X, atraj.UniformTimeTrajectoriesData
         ), "X must be uniform time"
-        X_, Xp_, U_ = X.next_step_matrices
+        if self.weights is None:
+            X_, Xp_, U_ = X.next_step_matrices
+            W = None
+        else:
+            X_, Xp_, U_, W = X.next_step_matrices_weights(self.weights)
         if self.normalize:
             X_ = self.scaler.fit_transform(X_.T).T
             Xp_ = self.scaler.transform(Xp_.T).T
-        self.fit_next_step(X_, Xp_, U_)
+        self.fit_next_step(X_, Xp_, U_, W)
         self.sampling_period = X.sampling_period
         self.names = X.state_names
 
@@ -75,18 +85,23 @@ class GradientEstimator(TrajectoryEstimator):
     """Estimator of discrete time dynamical systems
 
     Requires that the data be uniform time
-    
+
     :param normalize: apply MinMax normalization to the fit data
     :param feature_range: range for MinMax scaler
     """
 
-    def __init__(self, normalize=True, feature_range=(-1, 1)) -> None:
+    def __init__(self, normalize=True, feature_range=(-1, 1), weights=None) -> None:
         super().__init__(normalize=normalize, feature_range=feature_range)
         self.names = None
+        self.weights = None
 
     @abc.abstractmethod
     def fit_gradient(
-        self, X: np.ndarray, Y: np.ndarray, U: Optional[np.ndarray] = None
+        self,
+        X: np.ndarray,
+        Y: np.ndarray,
+        U: Optional[np.ndarray] = None,
+        weights: Optional[np.ndarray] = None,
     ) -> None:
         """an alternative fit method that uses a trajectories data structure"""
         pass
@@ -95,10 +110,14 @@ def fit(self, X: atraj.TrajectoriesData) -> None:
         assert isinstance(
             X, atraj.UniformTimeTrajectoriesData
         ), "X must be uniform time"
-        X_, Xp_, U_ = X.differentiate
+        if self.weights is None:
+            X_, Xp_, U_ = X.differentiate
+            W = None
+        else:
+            X_, Xp_, U_, W = X.differentiate_weights(self.weights)
         if self.normalize:
             X_ = self.scaler.fit_transform(X_.T).T
             Xp_ = self.scaler.transform(Xp_.T).T
-        self.fit_gradient(X_, Xp_, U_)
+        self.fit_gradient(X_, Xp_, U_, W)
         self.sampling_period = X.sampling_period
         self.names = X.state_names
diff --git a/autokoopman/core/scoring.py b/autokoopman/core/scoring.py
index 894a6d1..dace94e 100644
--- a/autokoopman/core/scoring.py
+++ b/autokoopman/core/scoring.py
@@ -30,7 +30,8 @@ def weighted_score(
 
         absdiff = (prediction_data - true_data).abs()
         end_errors = np.array(
-            [norm(weights_f[n] * s.states, axis=1) for n, s in absdiff._trajs.items()]
+            [norm(weights_f[n] * s.states, axis=1) for n, s in absdiff._trajs.items()],
+            dtype=object,
         )
         return np.sum(np.concatenate(end_errors, axis=0))
 
@@ -43,7 +44,7 @@ def end_point_score(true_data: TrajectoriesData, prediction_data: TrajectoriesDa
     @staticmethod
     def total_score(true_data: TrajectoriesData, prediction_data: TrajectoriesData):
         errors = (prediction_data - true_data).norm()
-        end_errors = np.array([s.states.flatten() for s in errors])
+        end_errors = np.array([s.states.flatten() for s in errors], dtype=object)
         return np.mean(np.concatenate(end_errors, axis=0))
 
     @staticmethod
diff --git a/autokoopman/core/trajectory.py b/autokoopman/core/trajectory.py
index 44d2372..3ebe55a 100644
--- a/autokoopman/core/trajectory.py
+++ b/autokoopman/core/trajectory.py
@@ -573,9 +573,7 @@ def from_pandas(
         return cls({k: v.to_uniform_time_traj() for k, v in traj_data._trajs.items()})
 
     @property
-    def next_step_matrices(
-        self, nstep=1
-    ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
+    def next_step_matrices(self) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
         r"""
         Next Step Snapshot Matrices
             Return the two "snapshot matrices" :math:`\mathbf X, \mathbf X'` of observations :math:`\{x_1, x_2, ..., x_n \}`,
@@ -601,23 +599,36 @@ def next_step_matrices(
             Brunton, S. L., & Kutz, J. N. (2022). Data-driven science and engineering: Machine learning,
             dynamical systems, and control. Cambridge University Press. pp 236-7
         """
-        return self.n_step_matrices(1)
+        return self.n_step_matrices(1, None)[:-1]
+
+    def next_step_matrices_weights(
+        self, weights, nstep=1
+    ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
+        r"""
+        Weighted Next Step Snapshot Matrices
+        """
+        return self.n_step_matrices(1, weights)
 
     def n_step_matrices(
-        self, nstep
+        self, nstep, weights=None
     ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
-        X = np.vstack([x.states[:-nstep:nstep, :] for _, x in self._trajs.items()]).T
+        items = self._trajs.items()
+        X = np.vstack([x.states[:-nstep:nstep, :] for _, x in items]).T
 
-        Xp = np.vstack([x.states[nstep::nstep, :] for _, x in self._trajs.items()]).T
+        Xp = np.vstack([x.states[nstep::nstep, :] for _, x in items]).T
 
         if self.input_names is not None:
-            U = np.vstack(
-                [u.inputs[:-nstep:nstep, :] for _, u in self._trajs.items()]
-            ).T
+            U = np.vstack([u.inputs[:-nstep:nstep, :] for _, u in items]).T
         else:
             U = None
 
-        return X, Xp, U
+        # collect weights
+        if weights is None:
+            W = None
+        else:
+            W = np.hstack([weights[idx].flatten()[1:] for idx, _ in items])
+
+        return X, Xp, U, W
 
     @property
     def differentiate(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
@@ -637,3 +648,28 @@ def differentiate(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
             U = None
 
         return X, Xp, U
+
+    def differentiate_weights(
+        self, weights
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        r"""
+        Numerical Differentiation with Weights -- Current State, Gradient State Estimate, Input, weight
+        """
+        items = list(self._trajs.items())
+
+        X = np.vstack([x.states[:-1, :] for _, x in items]).T
+
+        # finite difference
+        Xp = np.vstack([x.states[1:, :] for _, x in items]).T
+        Xp -= X
+        Xp /= self.sampling_period
+
+        if self.input_names is not None:
+            U = np.vstack([u.inputs[:-1, :] for _, u in items]).T
+        else:
+            U = None
+
+        # collect weights
+        W = np.hstack([weights[idx].flatten()[:-1] for idx, _ in items])
+
+        return X, Xp, U, W
diff --git a/autokoopman/estimator/deepkoopman.py b/autokoopman/estimator/deepkoopman.py
index c6077bc..5929f79 100644
--- a/autokoopman/estimator/deepkoopman.py
+++ b/autokoopman/estimator/deepkoopman.py
@@ -329,7 +329,7 @@ def _get_loss(mseloss, l1loss, x, xn, y, yn, xr, Xn):
             for step_size in range(1, self.rollout_steps + 1):
                 # upload to GPU and normalize the data
                 # NOTE: this is expensive :(
-                X, Xn, U = trajs.n_step_matrices(step_size)
+                X, Xn, U, _ = trajs.n_step_matrices(step_size)
                 X = torch.tensor(X.T, dtype=torch.float32).to(self.device)
                 Xn = torch.tensor(Xn.T, dtype=torch.float32).to(self.device)
                 U = (
diff --git a/autokoopman/estimator/koopman.py b/autokoopman/estimator/koopman.py
index 11e86c8..60d2981 100644
--- a/autokoopman/estimator/koopman.py
+++ b/autokoopman/estimator/koopman.py
@@ -30,6 +30,28 @@ def dmdc(X, Xp, U, r):
     return Atilde[:, :state_size], Atilde[:, state_size:]
 
 
+def wdmdc(X, Xp, U, r, W):
+    """Weighted Dynamic Mode Decomposition with Control (wDMDC)"""
+    # we allow weights to be either in diag or full representation
+    W = np.diag(W) if len(np.array(W).shape) == 2 else W
+
+    if U is not None:
+        Y = np.hstack((X, U)).T
+    else:
+        Y = X.T
+    Yp = Xp.T
+    Y, Yp = W * Y, W * Yp
+    state_size = Yp.shape[0]
+
+    # compute Atilde
+    U, Sigma, V = np.linalg.svd(Y, False)
+    U, Sigma, V = U[:, :r], np.diag(Sigma)[:r, :r], V.conj().T[:, :r]
+
+    # get the transformation
+    Atilde = Yp @ V @ np.linalg.inv(Sigma) @ U.conj().T
+    return Atilde[:, :state_size], Atilde[:, state_size:]
+
+
 class KoopmanDiscEstimator(kest.NextStepEstimator):
     """Koopman Discrete Estimator
 
@@ -43,6 +65,7 @@ class KoopmanDiscEstimator(kest.NextStepEstimator):
     :param sampling_period: sampling period of the uniform time, discrete system
     :param dim: dimension of the system state
     :param rank: rank of the DMDc
+    :param weights: observation weights (optional)
 
     References
         Proctor, J. L., Brunton, S. L., & Kutz, J. N. (2018). Generalizing Koopman theory to allow for inputs and control. SIAM Journal on Applied Dynamical Systems, 17(1), 909-930.
@@ -50,15 +73,19 @@ class KoopmanDiscEstimator(kest.NextStepEstimator):
         See https://epubs.siam.org/doi/pdf/10.1137/16M1062296 for more details
     """
 
-    def __init__(self, observables, sampling_period, dim, rank, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(self, observables, sampling_period, dim, rank, weights=None, **kwargs):
+        super().__init__(weights=weights, **kwargs)
         self.dim = dim
         self.obs = observables
         self.rank = int(rank)
         self._A, self._B = None, None
 
     def fit_next_step(
-        self, X: np.ndarray, Y: np.ndarray, U: Optional[np.ndarray] = None
+        self,
+        X: np.ndarray,
+        Y: np.ndarray,
+        U: Optional[np.ndarray] = None,
+        weights: Optional[np.ndarray] = None,
     ) -> None:
         """fits the discrete system model
 
@@ -66,7 +93,12 @@ def fit_next_step(
         """
         G = np.array([self.obs(xi).flatten() for xi in X.T])
         Gp = np.array([self.obs(xi).flatten() for xi in Y.T])
-        self._A, self._B = dmdc(G, Gp, U.T if U is not None else U, self.rank)
+        if weights is None:
+            self._A, self._B = dmdc(G, Gp, U.T if U is not None else U, self.rank)
+        else:
+            self._A, self._B = wdmdc(
+                G, Gp, U.T if U is not None else U, self.rank, weights
+            )
         self._has_input = U is not None
 
     @property
@@ -85,6 +117,7 @@ class KoopmanContinuousEstimator(kest.GradientEstimator):
     :param observables: function that returns the observables of the system state
     :param dim: dimension of the system state
     :param rank: rank of the DMDc
+    :param weights: observation weights (optional)
 
     References
         Proctor, J. L., Brunton, S. L., & Kutz, J. N. (2018). Generalizing Koopman theory to allow for inputs and control. SIAM Journal on Applied Dynamical Systems, 17(1), 909-930.
@@ -93,15 +126,19 @@ class KoopmanContinuousEstimator(kest.GradientEstimator):
 
     """
 
-    def __init__(self, observables, dim, rank, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(self, observables, dim, rank, weights=None, **kwargs):
+        super().__init__(weights=weights, **kwargs)
         self.dim = dim
         self.obs = observables
         self.rank = int(rank)
         self._A, self._B = None, None
 
     def fit_gradient(
-        self, X: np.ndarray, Y: np.ndarray, U: Optional[np.ndarray] = None
+        self,
+        X: np.ndarray,
+        Y: np.ndarray,
+        U: Optional[np.ndarray] = None,
+        weights: Optional[np.ndarray] = None,
     ) -> None:
         """fits the gradient system model
 
@@ -109,7 +146,12 @@ def fit_gradient(
         """
         G = np.array([self.obs(xi).flatten() for xi in X.T])
         Gp = np.array([self.obs(xi).flatten() for xi in Y.T])
-        self._A, self._B = dmdc(G, Gp, U.T if U is not None else U, self.rank)
+        if weights is None:
+            self._A, self._B = dmdc(G, Gp, U.T if U is not None else U, self.rank)
+        else:
+            self._A, self._B = wdmdc(
+                G, Gp, U.T if U is not None else U, self.rank, weights
+            )
         self._has_input = U is not None
 
     @property
diff --git a/autokoopman/tuner/gridsearch.py b/autokoopman/tuner/gridsearch.py
index e9081fd..29fb97f 100644
--- a/autokoopman/tuner/gridsearch.py
+++ b/autokoopman/tuner/gridsearch.py
@@ -66,5 +66,6 @@ def tune(
                 break
             except Exception as exc:
                 print(f"Error: {exc}")
+                raise exc
                 self.error_messages.append((param, exc))
         return self.best_result
diff --git a/notebooks/Koopman-Symbolic.ipynb b/notebooks/Koopman-Symbolic.ipynb
new file mode 100644
index 0000000..622da2f
--- /dev/null
+++ b/notebooks/Koopman-Symbolic.ipynb
@@ -0,0 +1,220 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "b9d1e95d-80e3-458c-a883-e080d0bf3a8e",
+   "metadata": {},
+   "source": [
+    "# Koopman Linearization of Symbolic Non-Linear Systems\n",
+    "\n",
+    "We seek to find Koopman observable for non-linear systems of the symbolic closed form\n",
+    "\n",
+    "$$\n",
+    "\\dot{\\mathbf x} = \\begin{bmatrix} f_1(\\mathbf x) \\\\ \\vdots \\\\ f_m(\\mathbf x) \\end{bmatrix}.\n",
+    "$$\n",
+    "\n",
+    "**Lie Derivative.** For a vector field $F = \\begin{bmatrix}f_1 \\\\ f_2 \\\\ \\vdots \\\\ f_m \\end{bmatrix}$, the *Lie Derivative* of a smooth function $f(\\mathbf x)$ is given by\n",
+    "\n",
+    "$$\n",
+    "\\mathcal L_F(f) = \\nabla f \\cdot F(\\mathbf x) = \\sum_{i=1}^n \\left( \\frac{\\partial f}{\\partial \\mathbf x_i} \\cdot f_i \\right).\n",
+    "$$\n",
+    "\n",
+    "## Automatic Abstraction of Symbolic Model [Sankaranarayanan, S. (2011, April)]\n",
+    "\n",
+    "1. Start with a sparse, symbolic model defined by the nonlinear equations $\\dot {\\mathbf x} = F(\\mathbf x)$. Choose an initial vector space $V_0 = \\operatorname{Span}(\\{f_1, \\cdots, f_k\\})$.\n",
+    "2. At $i$th iteration, refine the basis such that $V_{i+1} \\subseteq V_i, \\quad V_{i+1} = \\{f \\in V_i | \\mathcal L_F (f) \\in V_i\\}$.\n",
+    "3. Terminate at iteration $n$ when $V_{n+1}=V_{n}$.\n",
+    "\n",
+    "## Data-Driven (eDMD)\n",
+    "\n",
+    "1. Start with a sparse, symbolic model defined by the nonlinear equations $\\dot {\\mathbf x} = F(\\mathbf x)$. Choose basis functions $B=\\{f_1, \\cdots, f_k \\}$.\n",
+    "2. Run eDMD on the fixed observables\n",
+    "$$\n",
+    "g(\\mathbf x) = \\begin{bmatrix} \\mathbf x \\\\ \\mathcal L_F(f_1) \\\\ \\vdots \\\\  \\mathcal L_F(f_k) \\end{bmatrix}.\n",
+    "$$\n",
+    "\n",
+    "## References\n",
+    "\n",
+    "Sankaranarayanan, S. (2011, April). Automatic abstraction of non-linear systems using change of bases transformations. In Proceedings of the 14th international conference on Hybrid systems: computation and control (pp. 143-152).\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f4f243e3-da0b-47e8-85c4-a9c49f3ae4ee",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "sys.path.append(\"..\")\n",
+    "\n",
+    "import autokoopman.benchmark.fhn as pfhn\n",
+    "import autokoopman as ak\n",
+    "import sympy as sp\n",
+    "\n",
+    "from itertools import product\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "%matplotlib inline\n",
+    "%config InlineBackend.figure_format='retina'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "aa258ddd-cd57-4f9a-8843-cced8edede12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def lie_derivative(f, F, variables):\n",
+    "    \"\"\"symbolic lie derivative\"\"\"\n",
+    "    return sp.expand(sum(sp.diff(f, xi)*F[i] for i, xi in enumerate(variables)))\n",
+    "    \n",
+    "def update_basis(F, basis, variables):\n",
+    "    \"\"\"abstraction method\"\"\" \n",
+    "    # generate the candidate function\n",
+    "    basis_coef = sp.symbols(\" \".join([f\"c{i}\" for i in range(0, len(basis))]))\n",
+    "    f = sum(ci*bi for ci, bi in zip(basis_coef, basis))\n",
+    "\n",
+    "    # lie derivative\n",
+    "    lf = lie_derivative(f, F, variables)\n",
+    "\n",
+    "    # remove terms that do not belong to span(*F)\n",
+    "    subs = []\n",
+    "    for term in lf.args:\n",
+    "        in_span = False\n",
+    "        for bf in basis:\n",
+    "            if not any([(term / bf).has(xi) for xi in variables]):\n",
+    "                in_span = True\n",
+    "        if not in_span:\n",
+    "            subs.append([term.has(ci) for ci in basis_coef].index(True))\n",
+    "    return [bi for i, bi in enumerate(basis) if i not in subs]\n",
+    "\n",
+    "def obs(F, basis, variables):\n",
+    "    \"\"\"lie derivative observables\"\"\"\n",
+    "    return variables + [lie_derivative(bf, F, variables) for bf in basis]\n",
+    "\n",
+    "def generate_monomials(variables, order):\n",
+    "    \"\"\"generate initial basis functions as monomials up to an order\"\"\"\n",
+    "    monomials = set()\n",
+    "    # Generate all combinations of powers including and up to the order for each variable\n",
+    "    for powers in product(range(order + 1), repeat=len(variables)):\n",
+    "        if sum(powers) <= order:\n",
+    "            monomial = sp.prod([var**power for var, power in zip(variables, powers)])\n",
+    "            monomials.add(monomial)\n",
+    "    return list(monomials)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "948204d9-e295-4e96-baf2-47791e571111",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# build lie observables for the FHN system\n",
+    "fhn = pfhn.FitzHughNagumo()\n",
+    "variables = fhn._variables[1:]\n",
+    "exprs = fhn._exprs\n",
+    "x0, x1 = variables\n",
+    "basis = generate_monomials(variables, 2)[1:]\n",
+    "lie_obs = ak.core.observables.SymbolicObservable(variables, obs(exprs, basis, variables))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9c05d7bc-5f71-406c-9c87-bf928737c875",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# learn the Koopman operator (eDMD)\n",
+    "training_data = fhn.solve_ivps(\n",
+    "    initial_states=np.random.uniform(low=-2.0, high=2.0, size=(30, 2)),\n",
+    "    tspan=[0.0, 1.0],\n",
+    "    sampling_period=0.1\n",
+    ")\n",
+    "\n",
+    "# learn model from data\n",
+    "experiment_results = ak.auto_koopman(\n",
+    "    training_data,          # list of trajectories\n",
+    "    sampling_period=0.1,    # sampling period of trajectory snapshots\n",
+    "    obs_type=lie_obs,         # use Random Fourier Features Observables\n",
+    "    opt=\"monte-carlo\",             # grid search to find best hyperparameters\n",
+    "    n_obs=200,              # maximum number of observables to try\n",
+    "    max_opt_iter=10,       # maximum number of optimization iterations\n",
+    "    grid_param_slices=5,   # for grid search, number of slices for each parameter\n",
+    "    n_splits=5,             # k-folds validation for tuning, helps stabilize the scoring\n",
+    "    normalize=False,\n",
+    "    rank=(1, 200, 40)       # rank range (start, stop, step) DMD hyperparameter\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f9ca2671-5dca-4a57-a833-540fb5825406",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# plot the results on an extrapolation test set\n",
+    "model = experiment_results['tuned_model']\n",
+    "\n",
+    "testing_data = fhn.solve_ivps(\n",
+    "    initial_states=np.random.uniform(low=-3.0, high=3.0, size=(30, 2)),\n",
+    "    tspan=[0.0, 1.0],\n",
+    "    sampling_period=0.1\n",
+    ")\n",
+    "\n",
+    "# simulate using the learned model\n",
+    "iv = [0.5, 0.1]\n",
+    "prediction_data = model.solve_ivps(\n",
+    "    initial_states=[t.states[0] for t in testing_data],\n",
+    "    tspan=(0.0, 1.0),\n",
+    "    sampling_period=0.1\n",
+    ")\n",
+    "\n",
+    "# plot the results\n",
+    "plt.figure(figsize=(8,8))\n",
+    "for idx, trajectory in enumerate(prediction_data):\n",
+    "    plt.plot(*trajectory.states.T, 'r', label='prediction' if idx == 0 else None)\n",
+    "for idx, true_trajectory in enumerate(testing_data):\n",
+    "    plt.plot(*true_trajectory.states.T, 'k', label='ground truth' if idx == 0 else None)\n",
+    "for idx, trajectory in enumerate(training_data):\n",
+    "    plt.plot(*trajectory.states.T, 'g', label='training data' if idx == 0 else None)\n",
+    "plt.legend()\n",
+    "plt.grid()\n",
+    "plt.title(\"Lie Derivative Observables Extrapolation Test\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0098f2e2-7c17-414e-85e1-5f3bfcade1d7",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/notebooks/weighted-cost-func.ipynb b/notebooks/weighted-cost-func.ipynb
index a81416a..3629052 100644
--- a/notebooks/weighted-cost-func.ipynb
+++ b/notebooks/weighted-cost-func.ipynb
@@ -52,28 +52,29 @@
    "outputs": [],
    "source": [
     "# create trajectories as numpy array and create a weights array\n",
+    "# NOTE: learning_weights does not allow you to weight state, but the observations!\n",
     "trajectories = []\n",
     "weights = []\n",
     "\n",
     "# create weights for every time point\n",
     "for idx, traj in enumerate(training_data):\n",
+    "    # good trajectory\n",
     "    trajectories.append(traj.states)\n",
-    "    \n",
-    "    # uniform weights\n",
-    "    #weights.append(np.ones(len(traj.states)) / len(traj.states) )\n",
     "\n",
-    "    # distance weights\n",
-    "    #weights.append(traj.norm().states / (np.max(traj.norm().states) * len(traj.states)) )\n",
+    "    # garbage trajectory\n",
+    "    trajectories.append(np.random.rand(*traj.states.shape))\n",
     "    \n",
-    "    # individual state weights\n",
-    "    #w = np.abs(traj.states) / 10.0\n",
-    "    #w[:, 0] = 0.0\n",
-    "    #w = np.ones(traj.states.shape) / 10.0\n",
-    "    w = traj.abs().states\n",
+    "\n",
+    "    # weight good trajectory by its 1 norm\n",
+    "    w = np.sum(traj.abs().states, axis=1)\n",
+    "    weights.append(w)\n",
+    "\n",
+    "    # weight garbage trajectory to zero\n",
+    "    w = np.zeros(len(traj.states))\n",
     "    weights.append(w)\n",
     "\n",
     "# you can also use a dict to name the trajectories if using TrajectoriesData (numpy arrays are named by their index number)\n",
-    "#weights = {idx: w for idx, w in enumerate(weights)}"
+    "weights = {idx: w for idx, w in enumerate(weights)}"
    ]
   },
   {
@@ -83,12 +84,37 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# learn model from data\n",
+    "# learn model from weighted data\n",
     "experiment_results = auto_koopman(\n",
     "    trajectories,          # list of trajectories\n",
     "    sampling_period=0.1,    # sampling period of trajectory snapshots\n",
     "    obs_type=\"rff\",         # use Random Fourier Features Observables\n",
     "    cost_func=\"weighted\",   # use \"weighted\" cost function\n",
+    "    learning_weights=weights, # weight the eDMD algorithm objectives\n",
+    "    scoring_weights=weights, # pass weights as required for cost_func=\"weighted\"\n",
+    "    opt=\"grid\",             # grid search to find best hyperparameters\n",
+    "    n_obs=200,              # maximum number of observables to try\n",
+    "    max_opt_iter=200,       # maximum number of optimization iterations\n",
+    "    grid_param_slices=5,   # for grid search, number of slices for each parameter\n",
+    "    n_splits=5,             # k-folds validation for tuning, helps stabilize the scoring\n",
+    "    rank=(1, 200, 40)       # rank range (start, stop, step) DMD hyperparameter\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3eea17e2-46a9-493a-976a-273d83dd8fa6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# learn unweighted model from data\n",
+    "experiment_results_unweighted = auto_koopman(\n",
+    "    trajectories,          # list of trajectories\n",
+    "    sampling_period=0.1,    # sampling period of trajectory snapshots\n",
+    "    obs_type=\"rff\",         # use Random Fourier Features Observables\n",
+    "    cost_func=\"weighted\",   # use \"weighted\" cost function\n",
+    "    learning_weights=None,  # don't use eDMD weighting\n",
     "    scoring_weights=weights, # pass weights as required for cost_func=\"weighted\"\n",
     "    opt=\"grid\",             # grid search to find best hyperparameters\n",
     "    n_obs=200,              # maximum number of observables to try\n",
@@ -119,6 +145,7 @@
    "source": [
     "# get the model from the experiment results\n",
     "model = experiment_results['tuned_model']\n",
+    "model_uw = experiment_results_unweighted['tuned_model']\n",
     "\n",
     "# simulate using the learned model\n",
     "iv = [0.5, 0.1]\n",
@@ -126,6 +153,11 @@
     "    initial_state=iv,\n",
     "    tspan=(0.0, 10.0),\n",
     "    sampling_period=0.1\n",
+    ")\n",
+    "trajectory_uw = model_uw.solve_ivp(\n",
+    "    initial_state=iv,\n",
+    "    tspan=(0.0, 10.0),\n",
+    "    sampling_period=0.1\n",
     ")"
    ]
   },
@@ -146,7 +178,8 @@
     "plt.figure(figsize=(10, 6))\n",
     "\n",
     "# plot the results\n",
-    "plt.plot(*trajectory.states.T, label='Trajectory Prediction')\n",
+    "plt.plot(*trajectory.states.T, label='Weighted Trajectory Prediction')\n",
+    "plt.plot(*trajectory_uw.states.T, label='Trajectory Prediction')\n",
     "plt.plot(*true_trajectory.states.T, label='Ground Truth')\n",
     "\n",
     "plt.xlabel(\"$x_1$\")\n",
diff --git a/test/unit_test/test_autokoopman.py b/test/unit_test/test_autokoopman.py
index 3bf1c63..4550879 100644
--- a/test/unit_test/test_autokoopman.py
+++ b/test/unit_test/test_autokoopman.py
@@ -30,7 +30,6 @@ def test_autokoopman(obs, opt, cost, normalize):
     sp = 0.1
     ivs = np.random.uniform(low=-2.0, high=2.0, size=(20, 2))
     t_ends = [np.random.random() + 1.0 for idx in range(len(ivs))]
-    lengths = [int(t_end // sp) for t_end in t_ends]
     training_data = UniformTimeTrajectoriesData(
         {
             idx: fhn.solve_ivp(
@@ -45,7 +44,7 @@ def test_autokoopman(obs, opt, cost, normalize):
     # produce scoring weights if the cost is weighted
     if cost == "weighted":
         scoring_weights = {
-            idx: np.ones((length,)) * 0.01 for idx, length in enumerate(lengths)
+            idx: np.ones(len(traj.states)) * 0.01 for idx, traj in training_data._trajs.items() 
         }
     else:
         scoring_weights = None
@@ -86,8 +85,7 @@ def test_autokoopman_np(obs, opt, cost, normalize):
     # given issue #29, let's make these differently sized
     sp = 0.1
     ivs = np.random.uniform(low=-2.0, high=2.0, size=(20, 2))
-    t_ends = [np.random.random() + 1.0 for idx in range(len(ivs))]
-    lengths = [int(t_end // sp) for t_end in t_ends]
+    t_ends = [1.5 for idx in range(len(ivs))]
     training_data = UniformTimeTrajectoriesData(
         {
             idx: fhn.solve_ivp(
@@ -102,7 +100,7 @@ def test_autokoopman_np(obs, opt, cost, normalize):
     # produce scoring weights if the cost is weighted
     if cost == "weighted":
         scoring_weights = {
-            idx: np.ones((length,)) * 0.01 for idx, length in enumerate(lengths)
+            idx: np.ones((len(t.states),)) * 0.01 for idx, t in training_data._trajs.items()
         }
     else:
         scoring_weights = None
@@ -110,7 +108,7 @@ def test_autokoopman_np(obs, opt, cost, normalize):
     # put into numpy arrays for testing
     training_data_np = []
     scoring_weights_np = []
-    for name in training_data.traj_names:
+    for name in range(len(ivs)):
         training_data_np.append(training_data[name].states)
         if scoring_weights is not None:
             scoring_weights_np.append(scoring_weights[name])