fix lazy loader

neuro_py/__init__.py

@@ -1,20 +1,4 @@
-from lazy_loader import attach as _attach
+import lazy_loader as lazy
 
-_proto_all_ = [
-    "behavior",
-    "ensemble",
-    "io",
-    "lfp",
-    "plotting",
-    "process",
-    "session",
-    "spikes",
-    "stats",
-    "tuning",
-]
-
-__getattr__, __dir__, __all__ = _attach(
-    __name__, submodules=_proto_all_
-)
-
-del _attach
+(__getattr__, __dir__, __all__) = lazy.attach_stub(__name__, __file__)
+del lazy

neuro_py/__init__.pyi

@@ -0,0 +1,27 @@
+
+__all__ = [
+    "behavior",
+    "ensemble",
+    "io",
+    "lfp",
+    "plotting",
+    "process",
+    "session",
+    "spikes",
+    "stats",
+    "tuning",
+]
+
+
+from . import (
+    behavior,
+    ensemble,
+    io,
+    lfp,
+    plotting,
+    process,
+    session,
+    spikes,
+    stats,
+    tuning,
+)

neuro_py/behavior/__init__.py

@@ -1,15 +1,4 @@
-from . import cheeseboard
-from . import get_trials
-from . import linear_positions
-from . import linearization_pipeline
-from . import kinematics
-from . import well_traversal_classification
+import lazy_loader as lazy
 
-__all__ = [
-    "cheeseboard",
-    "get_trials",
-    "linear_positions",
-    "linearization_pipeline",
-    "kinematics",
-    "well_traversal_classification",
-]
+(__getattr__, __dir__, __all__) = lazy.attach_stub(__name__, __file__)
+del lazy

neuro_py/behavior/__init__.pyi

@@ -0,0 +1,54 @@
+__all__ = [
+    "plot_grid_with_circle_and_random_dots",
+    "get_linear_maze_trials",
+    "get_t_maze_trials",
+    "get_w_maze_trials",
+    "get_cheeseboard_trials",
+    "get_openfield_trials",
+    "get_velocity",
+    "get_speed",
+    "linearize_position",
+    "find_laps",
+    "peakdetz",
+    "find_good_laps",
+    "get_linear_track_lap_epochs",
+    "find_good_lap_epochs",
+    "NodePicker",
+    "paired_distances",
+    "enter_exit_target",
+    "enter_exit_target_dio",
+    "shift_well_enters",
+    "segment_path",
+    "find_last_non_center_well",
+    "get_correct_inbound_outbound",
+    "score_inbound_outbound",
+]
+
+from .cheeseboard import plot_grid_with_circle_and_random_dots
+from .get_trials import (
+    get_cheeseboard_trials,
+    get_linear_maze_trials,
+    get_openfield_trials,
+    get_t_maze_trials,
+    get_w_maze_trials,
+)
+from .kinematics import get_speed, get_velocity
+from .linear_positions import (
+    find_good_lap_epochs,
+    find_good_laps,
+    find_laps,
+    get_linear_track_lap_epochs,
+    linearize_position,
+    peakdetz,
+)
+from .linearization_pipeline import NodePicker
+from .well_traversal_classification import (
+    enter_exit_target,
+    enter_exit_target_dio,
+    find_last_non_center_well,
+    get_correct_inbound_outbound,
+    paired_distances,
+    score_inbound_outbound,
+    segment_path,
+    shift_well_enters,
+)

neuro_py/behavior/cheeseboard.py

@@ -1,14 +1,6 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
-from lazy_loader import attach as _attach
-
-__all__ = (
-    "plot_grid_with_circle_and_random_dots",
-)
-__getattr__, __dir__, __all__ = _attach(f"{__name__}", submodules=__all__)
-del _attach
-
 
 def plot_grid_with_circle_and_random_dots():
     # Create a 15x15 grid of dots within the circle
@@ -56,5 +48,6 @@ def plot_grid_with_circle_and_random_dots():
     plt.axis("off")
     plt.show()
 
+
 if __name__ == "__main__":
     plot_grid_with_circle_and_random_dots()

neuro_py/behavior/get_trials.py

@@ -5,23 +5,12 @@
 import nelpy as nel
 import numpy as np
 import scipy.io as sio
-from lazy_loader import attach as _attach
 from scipy.signal import medfilt
 
 from neuro_py.behavior import linear_positions, well_traversal_classification
 from neuro_py.io import loading
 from neuro_py.process.intervals import find_interval
 
-__all__ = (
-    "get_linear_maze_trials",
-    "get_t_maze_trials",
-    "get_w_maze_trials",
-    "get_cheeseboard_trials",
-    "get_openfield_trials",
-)
-__getattr__, __dir__, __all__ = _attach(f"{__name__}", submodules=__all__)
-del _attach
-
 
 # linear track
 def get_linear_maze_trials(basepath, epoch_input=None):

neuro_py/behavior/kinematics.py

@@ -2,22 +2,15 @@
 
 from typing import Union
 
-from lazy_loader import attach as _attach
 
-__all__ = (
-    "get_velocity",
-    "get_speed",
-)
-__getattr__, __dir__, __all__ = _attach(f"{__name__}", submodules=__all__)
-del _attach
-
-
-def get_velocity(position:np.ndarray, time:Union[np.ndarray,None]=None) -> np.ndarray:
+def get_velocity(
+    position: np.ndarray, time: Union[np.ndarray, None] = None
+) -> np.ndarray:
     if time is None:
         time = np.arange(position.shape[0])
     return np.gradient(position, time, axis=0)
 
 
-def get_speed(position:np.ndarray, time:Union[np.ndarray,None]=None) -> np.ndarray:
+def get_speed(position: np.ndarray, time: Union[np.ndarray, None] = None) -> np.ndarray:
     velocity = get_velocity(position, time=time)
     return np.sqrt(np.sum(velocity**2, axis=1))

neuro_py/behavior/linear_positions.py

@@ -3,20 +3,8 @@
 import nelpy as nel
 import numpy as np
 import pandas as pd
-from lazy_loader import attach as _attach
 from sklearn.decomposition import PCA
 
-__all__ = (
-    "linearize_position",
-    "find_laps",
-    "peakdetz",
-    "find_good_laps",
-    "get_linear_track_lap_epochs",
-    "find_good_lap_epochs",
-)
-__getattr__, __dir__, __all__ = _attach(f"{__name__}", submodules=__all__)
-del _attach
-
 
 def linearize_position(x, y):
     """

neuro_py/behavior/linearization_pipeline.py

@@ -6,18 +6,9 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
-from lazy_loader import attach as _attach
 from scipy.io import loadmat, savemat
 from track_linearization import get_linearized_position, make_track_graph
 
-__all__ = (
-    "NodePicker",
-    "load_animal_behavior",
-    "load_epoch",
-    "run",
-)
-__getattr__, __dir__, __all__ = _attach(f"{__name__}", submodules=__all__)
-
 
 """
 TODO: 

neuro_py/behavior/well_traversal_classification.py

@@ -2,24 +2,8 @@
 import numpy as np
 import pandas as pd
 
-from lazy_loader import attach as _attach
 from scipy.ndimage.measurements import label
 
-__all__ = (
-    "paired_distances",
-    "enter_exit_target",
-    "enter_exit_target_dio",
-    "shift_well_enters",
-    "segment_path",
-    "find_last_non_center_well",
-    "get_correct_inbound_outbound",
-    "score_inbound_outbound",
-)
-__getattr__, __dir__, __all__ = _attach(f"{__name__}", submodules=__all__)
-del _attach
-
-_WELL_NAMES = {1: "Center", 2: "Left", 3: "Right"}
-
 
 def paired_distances(x, y):
     """Euclidean distance between x and y at each time point.
@@ -88,9 +72,7 @@ def shift_well_enters(enter_exit):
     return shifted_enter_exit
 
 
-def segment_path(
-    time, position, well_locations, max_distance_from_well=10
-):
+def segment_path(time, position, well_locations, max_distance_from_well=10):
     """Label traversals between each well location.
 
     Parameters
@@ -206,7 +188,9 @@ def get_correct_inbound_outbound(segments_df):
 
 
 def score_inbound_outbound(
-    segments_df, min_distance_traveled=50, well_names=_WELL_NAMES
+    segments_df,
+    min_distance_traveled=50,
+    well_names={1: "Center", 2: "Left", 3: "Right"},
 ):
     """In the alternating arm task, determines whether the trial should be
     inbound (running to the center arm) or outbound (running to the opposite

neuro_py/ensemble/__init__.py

@@ -1,13 +1,4 @@
-from . import assembly_reactivation
-from . import assembly
-from . import explained_variance
-from . import similarity_index
-from . import similaritymat
+import lazy_loader as lazy
 
-__all__ = [
-    "assembly_reactivation",
-    "assembly",
-    "explained_variance",
-    "similarity_index",
-    "similaritymat",
-]
+(__getattr__, __dir__, __all__) = lazy.attach_stub(__name__, __file__)
+del lazy

neuro_py/ensemble/__init__.pyi

@@ -0,0 +1,31 @@
+__all__ = [
+    "toyExample",
+    "marcenkopastur",
+    "getlambdacontrol",
+    "binshuffling",
+    "circshuffling",
+    "runSignificance",
+    "extractPatterns",
+    "runPatterns",
+    "computeAssemblyActivity",
+    "AssemblyReact",
+    "ExplainedVariance",
+    "similarity_index",
+    "similaritymat",
+]
+
+from .assembly import (
+    binshuffling,
+    circshuffling,
+    computeAssemblyActivity,
+    extractPatterns,
+    getlambdacontrol,
+    marcenkopastur,
+    runPatterns,
+    runSignificance,
+    toyExample,
+)
+from .assembly_reactivation import AssemblyReact
+from .explained_variance import ExplainedVariance
+from .similarity_index import similarity_index
+from .similaritymat import similaritymat

neuro_py/ensemble/assembly.py

@@ -9,24 +9,9 @@
 from typing import Tuple, Union
 
 import numpy as np
-from lazy_loader import attach as _attach
 from scipy import stats
 from sklearn.decomposition import PCA, FastICA
 
-__all__ = (
-    "toyExample",
-    "ToyAssemblies",
-    "marcenkopastur",
-    "getlambdacontrol",
-    "binshuffling",
-    "circshuffling",
-    "runSignificance",
-    "extractPatterns",
-    "runPatterns",
-    "computeAssemblyActivity",
-)
-__getattr__, __dir__, __all__ = _attach(f"{__name__}", submodules=__all__)
-del _attach
 
 __author__ = "Vítor Lopes dos Santos"
 __version__ = "2019.1"