Skip to content

7. traj_recorder

torchfsm.traj_recorder.IntervalController ¤

A class to control the recording of trajectories at specified intervals. This class can be used as an input for the control_func parameters of recorder objects.

Parameters:

Name Type Description Default
interval int

The interval at which to record the trajectory.

 1
start int

The step at which to start recording the trajectory.

 0
Source code in torchfsm/traj_recorder.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class IntervalController():
    """
    A class to control the recording of trajectories at specified intervals.
        This class can be used as an input for the `control_func` parameters of recorder objects.

    Args:
        interval (int): The interval at which to record the trajectory.
        start (int): The step at which to start recording the trajectory.
    """

    def __init__(self,interval:int=1,start:int=0,) -> None:
        self.start=start
        self.interval=interval

    def __call__(self, step:int) -> bool:
        return step>=self.start and (step-self.start)%self.interval==0
start instance-attribute ¤
start = start
interval instance-attribute ¤
interval = interval
__init__ ¤
__init__(interval: int = 1, start: int = 0) -> None
Source code in torchfsm/traj_recorder.py 
17
18
19
def __init__(self,interval:int=1,start:int=0,) -> None:
    self.start=start
    self.interval=interval
__call__ ¤
__call__(step: int) -> bool
Source code in torchfsm/traj_recorder.py 
21
22
def __call__(self, step:int) -> bool:
    return step>=self.start and (step-self.start)%self.interval==0

torchfsm.traj_recorder._TrajRecorder ¤

A base class for trajectory recorders. A recorder is an object that helps to control the recording of trajectories during a simulation

Parameters:

Name Type Description Default
control_func Optional[Callable[[int], bool]]

A function that takes a step as input and returns a boolean indicating whether to record the trajectory at that step.

 None
include_initial_state bool

If True, the initial state will be included in the trajectory.

 True
Source code in torchfsm/traj_recorder.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
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
class _TrajRecorder():

    """
    A base class for trajectory recorders.
        A recorder is an object that helps to control the recording of trajectories during a simulation

    Args:
        control_func (Optional[Callable[[int],bool]]): A function that takes a step as input and returns a boolean indicating whether to record the trajectory at that step.
        include_initial_state (bool): If True, the initial state will be included in the trajectory.
    """

    def __init__(self,
                 control_func:Optional[Callable[[int],bool]]=None,
                 include_initial_state:bool=True,
                 ):
        control_func=default(control_func,lambda step: True)
        if include_initial_state:
            self.control_func=control_func
        else:
            self.control_func=lambda step: False if step==0 else control_func(step)
        self.return_in_fourier=False

    def record(self,step:int,frame:torch.tensor):
        """
        Record the trajectory at a given step.

        Args:
            step (int): The current step.
            frame (torch.tensor): The current frame to be recorded.
        """
        if self.control_func(step):
            self._record(step,frame)

    def _record(self,step:int,frame:torch.tensor):
        """
        Record the trajectory at a given step.
            This method should be implemented by subclasses.

        Args:
            step (int): The current step.
            frame (torch.tensor): The current frame to be recorded.
        """
        raise NotImplementedError

    def _traj_ifft(self,trajectory:torch.tensor):
        """
        Perform an inverse FFT on the trajectory.

        Args:
            trajectory (torch.tensor): The trajectory to be transformed.

        Returns:
            torch.tensor: The transformed trajectory.
        """
        fft_dim=tuple(-1*(i+1) for i in range(len(trajectory.shape)-3))
        return torch.fft.ifftn(trajectory,dim=fft_dim)

    @property
    def trajectory(self):
        """
        Get the recorded trajectory.
            This method should be implemented by subclasses.

        Args:
            return_in_fourier (bool): If True, return the trajectory in Fourier space. Default is False.

        Returns:
            torch.tensor: The recorded trajectory.
        """
        raise NotImplementedError
control_func instance-attribute ¤
control_func = control_func
return_in_fourier instance-attribute ¤
return_in_fourier = False
trajectory property ¤
trajectory

Get the recorded trajectory. This method should be implemented by subclasses.

Parameters:

Name Type Description Default
return_in_fourier bool

If True, return the trajectory in Fourier space. Default is False.

 required

Returns:

Type Description

torch.tensor: The recorded trajectory.
__init__ ¤
__init__(
    control_func: Optional[Callable[[int], bool]] = None,
    include_initial_state: bool = True,
)
Source code in torchfsm/traj_recorder.py 
35
36
37
38
39
40
41
42
43
44
def __init__(self,
             control_func:Optional[Callable[[int],bool]]=None,
             include_initial_state:bool=True,
             ):
    control_func=default(control_func,lambda step: True)
    if include_initial_state:
        self.control_func=control_func
    else:
        self.control_func=lambda step: False if step==0 else control_func(step)
    self.return_in_fourier=False
record ¤
record(step: int, frame: torch.tensor)

Record the trajectory at a given step.

Parameters:

Name Type Description Default
step int

The current step.

 required
frame tensor

The current frame to be recorded.

 required
Source code in torchfsm/traj_recorder.py 
46
47
48
49
50
51
52
53
54
55
def record(self,step:int,frame:torch.tensor):
    """
    Record the trajectory at a given step.

    Args:
        step (int): The current step.
        frame (torch.tensor): The current frame to be recorded.
    """
    if self.control_func(step):
        self._record(step,frame)
_record ¤
_record(step: int, frame: torch.tensor)

Record the trajectory at a given step. This method should be implemented by subclasses.

Parameters:

Name Type Description Default
step int

The current step.

 required
frame tensor

The current frame to be recorded.

 required
Source code in torchfsm/traj_recorder.py 
57
58
59
60
61
62
63
64
65
66
def _record(self,step:int,frame:torch.tensor):
    """
    Record the trajectory at a given step.
        This method should be implemented by subclasses.

    Args:
        step (int): The current step.
        frame (torch.tensor): The current frame to be recorded.
    """
    raise NotImplementedError
_traj_ifft ¤
_traj_ifft(trajectory: torch.tensor)

Perform an inverse FFT on the trajectory.

Parameters:

Name Type Description Default
trajectory tensor

The trajectory to be transformed.

 required

Returns:

Type Description

torch.tensor: The transformed trajectory.
Source code in torchfsm/traj_recorder.py 
68
69
70
71
72
73
74
75
76
77
78
79
def _traj_ifft(self,trajectory:torch.tensor):
    """
    Perform an inverse FFT on the trajectory.

    Args:
        trajectory (torch.tensor): The trajectory to be transformed.

    Returns:
        torch.tensor: The transformed trajectory.
    """
    fft_dim=tuple(-1*(i+1) for i in range(len(trajectory.shape)-3))
    return torch.fft.ifftn(trajectory,dim=fft_dim)

torchfsm.traj_recorder.AutoRecorder ¤

Bases: _TrajRecorder

A recorder that save the trajectory at the same devices as the simulation.

Parameters:

Name Type Description Default
control_func Optional[Callable[[int], bool]]

A function that takes a step as input and returns a boolean indicating whether to record the trajectory at that step.

 None
include_initial_state bool

If True, the initial state will be included in the trajectory.

 True
Source code in torchfsm/traj_recorder.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
class AutoRecorder(_TrajRecorder):

    """
    A recorder that save the trajectory at the same devices as the simulation.

    Args:
        control_func (Optional[Callable[[int],bool]]): A function that takes a step as input and returns a boolean indicating whether to record the trajectory at that step.
        include_initial_state (bool): If True, the initial state will be included in the trajectory.
    """

    def __init__(self,
                 control_func:Optional[Callable[[int],bool]]=None,
                 include_initial_state:bool=True,
                 ):
        super().__init__(control_func,include_initial_state)
        self._trajectory=[]

    def _record(self,step:int,frame:torch.tensor):
        if not isinstance(self._trajectory,torch.Tensor):
            self._trajectory.append(copy.deepcopy(frame))
        else:
            raise RuntimeError("The trajectory has been finalized.")

    @property
    def trajectory(self):
        if len(self._trajectory)==0:
            return None
        if self.return_in_fourier:
            return torch.stack(self._trajectory,dim=1)
        else:
            return self._traj_ifft(torch.stack(self._trajectory,dim=1)).real
control_func instance-attribute ¤
control_func = control_func
return_in_fourier instance-attribute ¤
return_in_fourier = False
_trajectory instance-attribute ¤
_trajectory = []
trajectory property ¤
trajectory
record ¤
record(step: int, frame: torch.tensor)

Record the trajectory at a given step.

Parameters:

Name Type Description Default
step int

The current step.

 required
frame tensor

The current frame to be recorded.

 required
Source code in torchfsm/traj_recorder.py 
46
47
48
49
50
51
52
53
54
55
def record(self,step:int,frame:torch.tensor):
    """
    Record the trajectory at a given step.

    Args:
        step (int): The current step.
        frame (torch.tensor): The current frame to be recorded.
    """
    if self.control_func(step):
        self._record(step,frame)
_traj_ifft ¤
_traj_ifft(trajectory: torch.tensor)

Perform an inverse FFT on the trajectory.

Parameters:

Name Type Description Default
trajectory tensor

The trajectory to be transformed.

 required

Returns:

Type Description

torch.tensor: The transformed trajectory.
Source code in torchfsm/traj_recorder.py 
68
69
70
71
72
73
74
75
76
77
78
79
def _traj_ifft(self,trajectory:torch.tensor):
    """
    Perform an inverse FFT on the trajectory.

    Args:
        trajectory (torch.tensor): The trajectory to be transformed.

    Returns:
        torch.tensor: The transformed trajectory.
    """
    fft_dim=tuple(-1*(i+1) for i in range(len(trajectory.shape)-3))
    return torch.fft.ifftn(trajectory,dim=fft_dim)
__init__ ¤
__init__(
    control_func: Optional[Callable[[int], bool]] = None,
    include_initial_state: bool = True,
)
Source code in torchfsm/traj_recorder.py 
105
106
107
108
109
110
def __init__(self,
             control_func:Optional[Callable[[int],bool]]=None,
             include_initial_state:bool=True,
             ):
    super().__init__(control_func,include_initial_state)
    self._trajectory=[]
_record ¤
_record(step: int, frame: torch.tensor)
Source code in torchfsm/traj_recorder.py 
112
113
114
115
116
def _record(self,step:int,frame:torch.tensor):
    if not isinstance(self._trajectory,torch.Tensor):
        self._trajectory.append(copy.deepcopy(frame))
    else:
        raise RuntimeError("The trajectory has been finalized.")

torchfsm.traj_recorder.CPURecorder ¤

Bases: AutoRecorder

A recorder that saves the trajectory on the CPU memory. This is useful for large trajectories that may not fit in GPU memory during simulation.

Parameters:

Name Type Description Default
control_func Optional[Callable[[int], bool]]

A function that takes a step as input and returns a boolean indicating whether to record the trajectory at that step.

 None
include_initial_state bool

If True, the initial state will be included in the trajectory.

 True
Source code in torchfsm/traj_recorder.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
class CPURecorder(AutoRecorder):

    """
    A recorder that saves the trajectory on the CPU memory.
        This is useful for large trajectories that may not fit in GPU memory during simulation.

    Args:
        control_func (Optional[Callable[[int],bool]]): A function that takes a step as input and returns a boolean indicating whether to record the trajectory at that step.
        include_initial_state (bool): If True, the initial state will be included in the trajectory.
    """

    def __init__(self,
                 control_func:Optional[Callable[[int],bool]]=None,
                 include_initial_state:bool=True
                 ):
        super().__init__(control_func,include_initial_state)

    def _record(self,step:int,frame:torch.tensor):
        if frame.is_cpu:
            self._trajectory.append(copy.deepcopy(frame))
        else:
            self._trajectory.append(frame.cpu())
control_func instance-attribute ¤
control_func = control_func
return_in_fourier instance-attribute ¤
return_in_fourier = False
trajectory property ¤
trajectory
_trajectory instance-attribute ¤
_trajectory = []
record ¤
record(step: int, frame: torch.tensor)

Record the trajectory at a given step.

Parameters:

Name Type Description Default
step int

The current step.

 required
frame tensor

The current frame to be recorded.

 required
Source code in torchfsm/traj_recorder.py 
46
47
48
49
50
51
52
53
54
55
def record(self,step:int,frame:torch.tensor):
    """
    Record the trajectory at a given step.

    Args:
        step (int): The current step.
        frame (torch.tensor): The current frame to be recorded.
    """
    if self.control_func(step):
        self._record(step,frame)
_traj_ifft ¤
_traj_ifft(trajectory: torch.tensor)

Perform an inverse FFT on the trajectory.

Parameters:

Name Type Description Default
trajectory tensor

The trajectory to be transformed.

 required

Returns:

Type Description

torch.tensor: The transformed trajectory.
Source code in torchfsm/traj_recorder.py 
68
69
70
71
72
73
74
75
76
77
78
79
def _traj_ifft(self,trajectory:torch.tensor):
    """
    Perform an inverse FFT on the trajectory.

    Args:
        trajectory (torch.tensor): The trajectory to be transformed.

    Returns:
        torch.tensor: The transformed trajectory.
    """
    fft_dim=tuple(-1*(i+1) for i in range(len(trajectory.shape)-3))
    return torch.fft.ifftn(trajectory,dim=fft_dim)
__init__ ¤
__init__(
    control_func: Optional[Callable[[int], bool]] = None,
    include_initial_state: bool = True,
)
Source code in torchfsm/traj_recorder.py 
138
139
140
141
142
def __init__(self,
             control_func:Optional[Callable[[int],bool]]=None,
             include_initial_state:bool=True
             ):
    super().__init__(control_func,include_initial_state)
_record ¤
_record(step: int, frame: torch.tensor)
Source code in torchfsm/traj_recorder.py 
144
145
146
147
148
def _record(self,step:int,frame:torch.tensor):
    if frame.is_cpu:
        self._trajectory.append(copy.deepcopy(frame))
    else:
        self._trajectory.append(frame.cpu())

torchfsm.traj_recorder.DiskRecorder ¤

Bases: _TrajRecorder

A recorder that saves the trajectory on the disk. This is useful for large trajectories that may not fit in GPU memory during simulation. The trajectory is saved in a temporary cache and then written to disk at specified intervals.

Parameters:

Name Type Description Default
control_func Optional[Callable[[int], bool]]

A function that takes a step as input and returns a boolean indicating whether to record the trajectory at that step.

 None
include_initial_state bool

If True, the initial state will be included in the trajectory.

 True
cache_dir Optional[str]

The directory where the trajectory will be saved. Default is "./saved_traj/".

 None
cache_freq int

The frequency at which to save the trajectory to disk. Default is 1.

 1
temp_cache_loc Literal['auto', 'cpu']

The location of the temporary cache. Default is "cpu".

 'cpu'
save_format Literal['numpy', 'torch']

The format in which to save the trajectory. Default is "torch".

 'torch'
Source code in torchfsm/traj_recorder.py 
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
class DiskRecorder(_TrajRecorder):

    """
    A recorder that saves the trajectory on the disk.
        This is useful for large trajectories that may not fit in GPU memory during simulation.
        The trajectory is saved in a temporary cache and then written to disk at specified intervals.

    Args:
        control_func (Optional[Callable[[int],bool]]): A function that takes a step as input and returns a boolean indicating whether to record the trajectory at that step.
        include_initial_state (bool): If True, the initial state will be included in the trajectory.
        cache_dir (Optional[str]): The directory where the trajectory will be saved. Default is "./saved_traj/".
        cache_freq (int): The frequency at which to save the trajectory to disk. Default is 1.
        temp_cache_loc (Literal["auto","cpu"]): The location of the temporary cache. Default is "cpu".
        save_format (Literal["numpy","torch"]): The format in which to save the trajectory. Default is "torch".
    """

    def __init__(self,
                 control_func:Optional[Callable[[int],bool]]=None,
                 include_initial_state:bool=True,
                 cache_dir:Optional[str]=None,
                 cache_freq:int=1,
                 temp_cache_loc:Literal["auto","cpu"]="cpu",
                 save_format:Literal["numpy","torch"]="torch",
                 ):
        super().__init__(control_func,include_initial_state)
        self.cache_dir=default(cache_dir,"./saved_traj/")
        self.cache_freq=cache_freq
        self._trajectory=[]
        self.temp_cache_loc=temp_cache_loc
        if self.temp_cache_loc not in ["auto","cpu"]:
            raise ValueError("temp_cache_loc must be either 'auto' or 'cpu'.")
        self.save_format=save_format
        if self.save_format not in ["numpy","torch"]:
            raise ValueError("save_format must be either 'numpy' or 'torch'.")

    def _record(self,step:int,frame:torch.tensor):
        if len(self.traj<self.cache_freq):    
            if self.temp_cache_loc=="cpu" and not frame.is_cpu:
                self._trajectory.append(frame.cpu())
            else:
                self._trajectory.append(copy.deepcopy(frame))
        else:
            temp_cache=torch.stack(self._trajectory,dim=1)
            temp_cache=temp_cache.to("cpu") if not temp_cache.is_cpu else temp_cache
            if not self.return_in_fourier:
                temp_cache=self._traj_ifft(temp_cache).real
            if self.save_format=="numpy":
                np.save(self.cache_dir+f"temp_cache_{step}",temp_cache.numpy()) 
            else:
                torch.save(temp_cache,self.cache_dir+f"temp_cache_{step}")     
            self._trajectory=[]      

    @property
    def trajectory(self):
        return None
control_func instance-attribute ¤
control_func = control_func
return_in_fourier instance-attribute ¤
return_in_fourier = False
cache_dir instance-attribute ¤
cache_dir = default(cache_dir, './saved_traj/')
cache_freq instance-attribute ¤
cache_freq = cache_freq
_trajectory instance-attribute ¤
_trajectory = []
temp_cache_loc instance-attribute ¤
temp_cache_loc = temp_cache_loc
save_format instance-attribute ¤
save_format = save_format
trajectory property ¤
trajectory
record ¤
record(step: int, frame: torch.tensor)

Record the trajectory at a given step.

Parameters:

Name Type Description Default
step int

The current step.

 required
frame tensor

The current frame to be recorded.

 required
Source code in torchfsm/traj_recorder.py 
46
47
48
49
50
51
52
53
54
55
def record(self,step:int,frame:torch.tensor):
    """
    Record the trajectory at a given step.

    Args:
        step (int): The current step.
        frame (torch.tensor): The current frame to be recorded.
    """
    if self.control_func(step):
        self._record(step,frame)
_traj_ifft ¤
_traj_ifft(trajectory: torch.tensor)

Perform an inverse FFT on the trajectory.

Parameters:

Name Type Description Default
trajectory tensor

The trajectory to be transformed.

 required

Returns:

Type Description

torch.tensor: The transformed trajectory.
Source code in torchfsm/traj_recorder.py 
68
69
70
71
72
73
74
75
76
77
78
79
def _traj_ifft(self,trajectory:torch.tensor):
    """
    Perform an inverse FFT on the trajectory.

    Args:
        trajectory (torch.tensor): The trajectory to be transformed.

    Returns:
        torch.tensor: The transformed trajectory.
    """
    fft_dim=tuple(-1*(i+1) for i in range(len(trajectory.shape)-3))
    return torch.fft.ifftn(trajectory,dim=fft_dim)
__init__ ¤
__init__(
    control_func: Optional[Callable[[int], bool]] = None,
    include_initial_state: bool = True,
    cache_dir: Optional[str] = None,
    cache_freq: int = 1,
    temp_cache_loc: Literal["auto", "cpu"] = "cpu",
    save_format: Literal["numpy", "torch"] = "torch",
)
Source code in torchfsm/traj_recorder.py 
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
def __init__(self,
             control_func:Optional[Callable[[int],bool]]=None,
             include_initial_state:bool=True,
             cache_dir:Optional[str]=None,
             cache_freq:int=1,
             temp_cache_loc:Literal["auto","cpu"]="cpu",
             save_format:Literal["numpy","torch"]="torch",
             ):
    super().__init__(control_func,include_initial_state)
    self.cache_dir=default(cache_dir,"./saved_traj/")
    self.cache_freq=cache_freq
    self._trajectory=[]
    self.temp_cache_loc=temp_cache_loc
    if self.temp_cache_loc not in ["auto","cpu"]:
        raise ValueError("temp_cache_loc must be either 'auto' or 'cpu'.")
    self.save_format=save_format
    if self.save_format not in ["numpy","torch"]:
        raise ValueError("save_format must be either 'numpy' or 'torch'.")
_record ¤
_record(step: int, frame: torch.tensor)
Source code in torchfsm/traj_recorder.py 
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
def _record(self,step:int,frame:torch.tensor):
    if len(self.traj<self.cache_freq):    
        if self.temp_cache_loc=="cpu" and not frame.is_cpu:
            self._trajectory.append(frame.cpu())
        else:
            self._trajectory.append(copy.deepcopy(frame))
    else:
        temp_cache=torch.stack(self._trajectory,dim=1)
        temp_cache=temp_cache.to("cpu") if not temp_cache.is_cpu else temp_cache
        if not self.return_in_fourier:
            temp_cache=self._traj_ifft(temp_cache).real
        if self.save_format=="numpy":
            np.save(self.cache_dir+f"temp_cache_{step}",temp_cache.numpy()) 
        else:
            torch.save(temp_cache,self.cache_dir+f"temp_cache_{step}")     
        self._trajectory=[]