Sawtooth wave function purely in pytorch

sawtooth_wave.py

import torch as th

"""
Author: Josue N Rivera
Date: 8/22/2024
Description: Function to generate sawtooth wave only using PyTorch
"""

def sawtooth_wave(t, period = 0.5, amplitude = 1.0, phase = 0.0, move = 0.0):
    x = (t + phase) / period
    return amplitude * (x - th.floor(x)) + move
    
if __name__ == '__main__':
  import matplotlib.pyplot as plt
  
  t = th.linspace(0, 1, 500)
  
  # Sawtooth wave with a period of 0.5
  plt.plot(t, sawtooth_wave(t))
    
  # Wave cut before period ended
  plt.plot(t, sawtooth_wave(t, period = 0.3))
  plt.plot(t, sawtooth_wave(t, phase = 0.1))
  
  # Wave with increased amplitude and displaced
  plt.plot(t, sawtooth_wave(t, amplitude = 2, move = -1.0))