Skip to content

main.py

LinearizationParameters

Bases: SpyglassMixin, Lookup

Choose whether to use an HMM to linearize position. This can help when the eucledian distances between separate arms are too close and the previous position has some information about which arm the animal is on.

Source code in src/spyglass/linearization/v1/main.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
@schema
class LinearizationParameters(SpyglassMixin, dj.Lookup):
    """Choose whether to use an HMM to linearize position. This can help when
    the eucledian distances between separate arms are too close and the previous
    position has some information about which arm the animal is on."""

    definition = """
    linearization_param_name : varchar(80)   # name for this set of parameters
    ---
    use_hmm = 0 : int   # use HMM to determine linearization
    # How much to prefer route distances between successive time points that are closer to the euclidean distance. Smaller numbers mean the route distance is more likely to be close to the euclidean distance.
    route_euclidean_distance_scaling = 1.0 : float
    sensor_std_dev = 5.0 : float   # Uncertainty of position sensor (in cm).
    # Biases the transition matrix to prefer the current track segment.
    diagonal_bias = 0.5 : float
    """

TrackGraph

Bases: SpyglassMixin, Manual

Graph representation of track representing the spatial environment. Used for linearizing position.

Source code in src/spyglass/linearization/v1/main.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
@schema
class TrackGraph(SpyglassMixin, dj.Manual):
    """Graph representation of track representing the spatial environment.
    Used for linearizing position."""

    definition = """
    track_graph_name : varchar(80)
    ----
    environment : varchar(80)  # Type of Environment
    node_positions : blob      # 2D position of nodes, (n_nodes, 2)
    edges: blob                # shape (n_edges, 2)
    linear_edge_order : blob   # order of edges in linear space, (n_edges, 2)
    linear_edge_spacing : blob # space btwn edges in linear space, (n_edges,)
    edge_map = NULL : blob     # Maps one edge to another before linearization
    """

    def get_networkx_track_graph(self, track_graph_parameters=None):
        """Get the track graph as a networkx graph."""
        if track_graph_parameters is None:
            track_graph_parameters = self.fetch1()
        return make_track_graph(
            node_positions=track_graph_parameters["node_positions"],
            edges=track_graph_parameters["edges"],
        )

    def plot_track_graph(self, ax=None, draw_edge_labels=False, **kwds):
        """Plot the track graph in 2D position space."""
        track_graph = self.get_networkx_track_graph(
            track_graph_parameters=self.fetch1()
        )
        plot_track_graph(
            track_graph, ax=ax, draw_edge_labels=draw_edge_labels, **kwds
        )

    def plot_track_graph_as_1D(
        self,
        ax=None,
        axis="x",
        other_axis_start=0.0,
        draw_edge_labels=False,
        node_size=300,
        node_color="#1f77b4",
    ):
        """Plot the track graph in 1D to see how the linearization is set up."""
        track_graph_parameters = self.fetch1()
        track_graph = self.get_networkx_track_graph(
            track_graph_parameters=track_graph_parameters
        )
        plot_graph_as_1D(
            track_graph,
            edge_order=track_graph_parameters["linear_edge_order"],
            edge_spacing=track_graph_parameters["linear_edge_spacing"],
            ax=ax,
            axis=axis,
            other_axis_start=other_axis_start,
            draw_edge_labels=draw_edge_labels,
            node_size=node_size,
            node_color=node_color,
        )

get_networkx_track_graph(track_graph_parameters=None)

Get the track graph as a networkx graph.

Source code in src/spyglass/linearization/v1/main.py 
55
56
57
58
59
60
61
62
def get_networkx_track_graph(self, track_graph_parameters=None):
    """Get the track graph as a networkx graph."""
    if track_graph_parameters is None:
        track_graph_parameters = self.fetch1()
    return make_track_graph(
        node_positions=track_graph_parameters["node_positions"],
        edges=track_graph_parameters["edges"],
    )

plot_track_graph(ax=None, draw_edge_labels=False, **kwds)

Plot the track graph in 2D position space.

Source code in src/spyglass/linearization/v1/main.py 
64
65
66
67
68
69
70
71
def plot_track_graph(self, ax=None, draw_edge_labels=False, **kwds):
    """Plot the track graph in 2D position space."""
    track_graph = self.get_networkx_track_graph(
        track_graph_parameters=self.fetch1()
    )
    plot_track_graph(
        track_graph, ax=ax, draw_edge_labels=draw_edge_labels, **kwds
    )

plot_track_graph_as_1D(ax=None, axis='x', other_axis_start=0.0, draw_edge_labels=False, node_size=300, node_color='#1f77b4')

Plot the track graph in 1D to see how the linearization is set up.

Source code in src/spyglass/linearization/v1/main.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def plot_track_graph_as_1D(
    self,
    ax=None,
    axis="x",
    other_axis_start=0.0,
    draw_edge_labels=False,
    node_size=300,
    node_color="#1f77b4",
):
    """Plot the track graph in 1D to see how the linearization is set up."""
    track_graph_parameters = self.fetch1()
    track_graph = self.get_networkx_track_graph(
        track_graph_parameters=track_graph_parameters
    )
    plot_graph_as_1D(
        track_graph,
        edge_order=track_graph_parameters["linear_edge_order"],
        edge_spacing=track_graph_parameters["linear_edge_spacing"],
        ax=ax,
        axis=axis,
        other_axis_start=other_axis_start,
        draw_edge_labels=draw_edge_labels,
        node_size=node_size,
        node_color=node_color,
    )

LinearizedPositionV1

Bases: SpyglassMixin, Computed

Linearized position for a given interval

Source code in src/spyglass/linearization/v1/main.py 
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
@schema
class LinearizedPositionV1(SpyglassMixin, dj.Computed):
    """Linearized position for a given interval"""

    definition = """
    -> LinearizationSelection
    ---
    -> AnalysisNwbfile
    linearized_position_object_id : varchar(40)
    """

    def make(self, key):
        """Populate LinearizedPositionV1 table with the linearized position.

        The linearized position is computed from the position data in the
        PositionOutput table. Parameters for linearization are specified in
        LinearizationParameters and the track graph is specified in TrackGraph.
        The linearization function is defined by the track_linearization
        package. The resulting linearized position is stored in an
        AnalysisNwbfile and added as an entry in the LinearizedPositionV1 and
        LinearizedPositionOutput (Merge) tables.
        """
        orig_key = copy.deepcopy(key)
        logger.info(f"Computing linear position for: {key}")

        # Fetch position data (Spyglass logic - always tested)
        position_nwb = PositionOutput().fetch_nwb(
            {"merge_id": key["pos_merge_id"]}
        )[0]
        key["analysis_file_name"] = AnalysisNwbfile().create(
            position_nwb["nwb_file_name"]
        )
        position = np.asarray(
            position_nwb["position"].get_spatial_series().data
        )
        time = np.asarray(
            position_nwb["position"].get_spatial_series().timestamps
        )

        linearization_parameters = (
            LinearizationParameters()
            & {"linearization_param_name": key["linearization_param_name"]}
        ).fetch1()
        track_graph_info = (
            TrackGraph() & {"track_graph_name": key["track_graph_name"]}
        ).fetch1()

        track_graph = (
            TrackGraph & {"track_graph_name": key["track_graph_name"]}
        ).get_networkx_track_graph()

        # External dependency - MOCKABLE in tests
        linear_position_df = self._compute_linearized_position(
            position=position,
            time=time,
            track_graph=track_graph,
            track_graph_info=track_graph_info,
            linearization_parameters=linearization_parameters,
        )

        # External I/O - MOCKABLE in tests
        key["linearized_position_object_id"] = self._save_linearization_results(
            linear_position_df=linear_position_df,
            analysis_file_name=key["analysis_file_name"],
            nwb_file_name=position_nwb["nwb_file_name"],
        )

        # Database operations (Spyglass logic - always tested)
        self.insert1(key)

        from spyglass.linearization.merge import LinearizedPositionOutput

        part_name = to_camel_case(self.table_name.split("__")[-1])

        LinearizedPositionOutput._merge_insert(
            [orig_key], part_name=part_name, skip_duplicates=True
        )

    def _compute_linearized_position(
        self,
        position,
        time,
        track_graph,
        track_graph_info,
        linearization_parameters,
    ):
        """Compute linearized position (external dependency).

        This method wraps the call to the track_linearization package,
        making it easy to mock in tests for faster execution.

        Parameters
        ----------
        position : np.ndarray
            2D position data (n_time, 2)
        time : np.ndarray
            Timestamps for position data
        track_graph : networkx.Graph
            Track graph from TrackGraph table
        track_graph_info : dict
            Track graph parameters (edges, spacing, etc.)
        linearization_parameters : dict
            Linearization parameters (HMM, scaling, etc.)

        Returns
        -------
        linear_position_df : pd.DataFrame
            Linearized position with time column
        """
        linear_position_df = get_linearized_position(
            position=position,
            track_graph=track_graph,
            edge_spacing=track_graph_info["linear_edge_spacing"],
            edge_order=track_graph_info["linear_edge_order"],
            use_HMM=linearization_parameters["use_hmm"],
            route_euclidean_distance_scaling=linearization_parameters[
                "route_euclidean_distance_scaling"
            ],
            sensor_std_dev=linearization_parameters["sensor_std_dev"],
            diagonal_bias=linearization_parameters["diagonal_bias"],
            edge_map=track_graph_info["edge_map"],
        )

        linear_position_df["time"] = time

        return linear_position_df

    def _save_linearization_results(
        self,
        linear_position_df,
        analysis_file_name,
        nwb_file_name,
    ):
        """Save linearization results to NWB file (external I/O).

        This method wraps all file I/O operations, making it easy to
        mock in tests to avoid filesystem dependencies.

        Parameters
        ----------
        linear_position_df : pd.DataFrame
            Linearized position dataframe
        analysis_file_name : str
            Name of analysis NWB file
        nwb_file_name : str
            Name of source NWB file

        Returns
        -------
        linearized_position_object_id : str
            Object ID for linearized position in NWB file
        """
        nwb_analysis_file = AnalysisNwbfile()

        linearized_position_object_id = nwb_analysis_file.add_nwb_object(
            analysis_file_name=analysis_file_name,
            nwb_object=linear_position_df,
        )

        nwb_analysis_file.add(
            nwb_file_name=nwb_file_name,
            analysis_file_name=analysis_file_name,
        )

        return linearized_position_object_id

    def fetch1_dataframe(self) -> DataFrame:
        """Fetch a single dataframe."""
        _ = self.ensure_single_entry()
        return self.fetch_nwb()[0]["linearized_position"].set_index("time")

make(key)

Populate LinearizedPositionV1 table with the linearized position.

The linearized position is computed from the position data in the PositionOutput table. Parameters for linearization are specified in LinearizationParameters and the track graph is specified in TrackGraph. The linearization function is defined by the track_linearization package. The resulting linearized position is stored in an AnalysisNwbfile and added as an entry in the LinearizedPositionV1 and LinearizedPositionOutput (Merge) tables.

Source code in src/spyglass/linearization/v1/main.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
def make(self, key):
    """Populate LinearizedPositionV1 table with the linearized position.

    The linearized position is computed from the position data in the
    PositionOutput table. Parameters for linearization are specified in
    LinearizationParameters and the track graph is specified in TrackGraph.
    The linearization function is defined by the track_linearization
    package. The resulting linearized position is stored in an
    AnalysisNwbfile and added as an entry in the LinearizedPositionV1 and
    LinearizedPositionOutput (Merge) tables.
    """
    orig_key = copy.deepcopy(key)
    logger.info(f"Computing linear position for: {key}")

    # Fetch position data (Spyglass logic - always tested)
    position_nwb = PositionOutput().fetch_nwb(
        {"merge_id": key["pos_merge_id"]}
    )[0]
    key["analysis_file_name"] = AnalysisNwbfile().create(
        position_nwb["nwb_file_name"]
    )
    position = np.asarray(
        position_nwb["position"].get_spatial_series().data
    )
    time = np.asarray(
        position_nwb["position"].get_spatial_series().timestamps
    )

    linearization_parameters = (
        LinearizationParameters()
        & {"linearization_param_name": key["linearization_param_name"]}
    ).fetch1()
    track_graph_info = (
        TrackGraph() & {"track_graph_name": key["track_graph_name"]}
    ).fetch1()

    track_graph = (
        TrackGraph & {"track_graph_name": key["track_graph_name"]}
    ).get_networkx_track_graph()

    # External dependency - MOCKABLE in tests
    linear_position_df = self._compute_linearized_position(
        position=position,
        time=time,
        track_graph=track_graph,
        track_graph_info=track_graph_info,
        linearization_parameters=linearization_parameters,
    )

    # External I/O - MOCKABLE in tests
    key["linearized_position_object_id"] = self._save_linearization_results(
        linear_position_df=linear_position_df,
        analysis_file_name=key["analysis_file_name"],
        nwb_file_name=position_nwb["nwb_file_name"],
    )

    # Database operations (Spyglass logic - always tested)
    self.insert1(key)

    from spyglass.linearization.merge import LinearizedPositionOutput

    part_name = to_camel_case(self.table_name.split("__")[-1])

    LinearizedPositionOutput._merge_insert(
        [orig_key], part_name=part_name, skip_duplicates=True
    )

fetch1_dataframe()

Fetch a single dataframe.

Source code in src/spyglass/linearization/v1/main.py 
276
277
278
279
def fetch1_dataframe(self) -> DataFrame:
    """Fetch a single dataframe."""
    _ = self.ensure_single_entry()
    return self.fetch_nwb()[0]["linearized_position"].set_index("time")