code: def test_graph_sampler(test_dataset):
      from pytorch_forecasting.data.samplers import (
          TimeSynchronizedBatchSampler,
      )
  
      class NeighborhoodSampler(TimeSynchronizedBatchSampler):
          def construct_batch_groups(self, groups):
              batch_size = self.batch_size
              self.batch_size = 1
              super().construct_batch_groups(groups)
              self.batch_size = batch_size
  
          def __iter__(self):
              if self.shuffle:
                  batch_samples = np.random.permutation(len(self))
              else:
                  batch_samples = np.arange(len(self))
  
              def distance_to_weights(dist):
                  return 1 / (1e-2 + np.power(dist * 5, 2))
  
              # for each point, sample the neighborhood
              # get groups associated with chosen sample
              data_groups = self.sampler.data_source.data[
                  "groups"
              ].float()
              n_groups = data_groups.size(1)  # number time series ids
              for idx in batch_samples:
                  name = self._group_index[
                      idx
                  ]  # time-synchronized group name
                  sub_group_idx = self._sub_group_index[idx]
                  selected_index = self._groups[name][sub_group_idx]
                  # select all other indices in same time group
                  indices = self.sampler.data_source.index.iloc[
                      self._groups[name]
                  ]
                  selected_pos = indices["index_start"].iloc[
                      sub_group_idx
                  ]
                  # remove selected sample
                  indices = indices[
                      lambda x: x["sequence_id"]
                      != indices["sequence_id"].iloc[sub_group_idx]
                  ]
                  # filter duplicate timeseries
                  # indices = indices.sort_values("sequence_length").drop_duplicates("sequence_id", keep="last")
  
                  # calculate distances for corresponding groups
                  group_distances = torch.cdist(
                      data_groups[[selected_pos]],
                      data_groups[indices["index_start"].to_numpy()],
                      p=0,
                  )[0].numpy()
                  # filter out all samples without group-link but not itself
                  connected_samples = group_distances < n_groups
                  relevant_indices = indices.index[connected_samples]
                  sample_weights = distance_to_weights(
                      group_distances[connected_samples]
                  )  # calculate weights for sampling neighborhood
  
                  # sample random subset of neighborhood
                  batch_size = min(
                      len(relevant_indices), self.batch_size - 1
                  )
                  batch_indices = [selected_index] + np.random.choice(
                      relevant_indices,
                      p=sample_weights / sample_weights.sum(),
                      replace=False,
                      size=batch_size,
                  ).tolist()
                  yield batch_indices
  
      dl = test_dataset.to_dataloader(
          batch_sampler=NeighborhoodSampler(
              SequentialSampler(test_dataset),
              batch_size=200,
              shuffle=True,
          )
      )
      for idx, a in enumerate(dl):
          print(a[0]["groups"].shape)
          if idx > 100:
              break
      print(a)
  
imports: from copy import deepcopy
  import itertools
  import pickle
  from typing import Dict
  import networkx
  import numpy
  import pandas
  import pytest
  from sklearn.preprocessing import StandardScaler
  import torch
  from torch.utils.data.sampler import SequentialSampler
  from pytorch_forecasting.data import EncoderNormalizer, GroupNormalizer, NaNLabelEncoder, TimeSeriesDataSet
  from pytorch_forecasting.data.encoders import MultiNormalizer, TorchNormalizer
  from pytorch_forecasting.data.examples import get_stallion_data
  from pytorch_forecasting.data.timeseries import _find_end_indices
  from pytorch_forecasting.utils import to_list
label_desc: Use image node_id 19 for calling a global function with example usage: test_graph_sampler(test_dataset) and returns: unknown