diff --git a/src/more_math/Parser/TensorEvalVisitor.py b/src/more_math/Parser/TensorEvalVisitor.py index 12f2318..2112b69 100644 --- a/src/more_math/Parser/TensorEvalVisitor.py +++ b/src/more_math/Parser/TensorEvalVisitor.py @@ -128,77 +128,70 @@ class TensorEvalVisitor(MathExprVisitor): def visitSigmoidFunc(self, ctx): return torch.sigmoid(self.visit(ctx.expr())) def visitAnglFunc(self, ctx): return torch.angle(self.visit(ctx.expr())) - def visitSfftFunc(self, ctx): - s = self.shape - - hop_length = 256 - n_fft = 512 - if len(s) < 4: - raise ValueError("Input tensor must have at least 4 dimensions for SFFT operation. You might be forgetting to to use inverse function before leaving node.") - self.shape=(s[0], - s[1], - s[3] * hop_length) +def visitSfftFunc(self, ctx): + s = self.shape + hop_length = 256 + n_fft = 512 + print ("SFFT input shape:", s) - shp = torch.zeros(self.shape) + if len(s) < 3: + raise ValueError("Input tensor must have at least 3 dims (B, C, T) for SFFT") - self.variables['T'] = getIndexTensorAlongDim(shp, 2) - self.variables['B'] = getIndexTensorAlongDim(shp, 0) - self.variables['C'] = getIndexTensorAlongDim(shp, 1) - self.variables['S'] = getIndexTensorAlongDim(shp, 2) - self.variables['R'] = torch.full_like(shp,self.variables['R'].flatten()[0].item()) + # Input: (B, C, T) + out_shape = (s[0], + s[1], + n_fft // 2 + 1, # frequency bins + (s[2] // hop_length) + 1) # frames + self.shape = out_shape - val = self.visit(ctx.expr()); - self.shape = s; + shp = torch.zeros(out_shape) + self.variables['T'] = getIndexTensorAlongDim(shp, 3) # time frames + self.variables['B'] = getIndexTensorAlongDim(shp, 0) + self.variables['C'] = getIndexTensorAlongDim(shp, 1) + self.variables['R'] = torch.full_like(shp, self.variables['R'].flatten()[0].item()) - shp = torch.zeros(self.shape) + val = self.visit(ctx.expr()) - vall = time_to_freq(val, n_fft, hop_length) - self.variables['T'] = getIndexTensorAlongDim(shp, 3) - self.variables['B'] = getIndexTensorAlongDim(shp, 0) - self.variables['C'] = getIndexTensorAlongDim(shp, 1) - self.variables['S'] = getIndexTensorAlongDim(shp, 2) - self.variables['R'] = torch.full_like(shp,self.variables['R'].flatten()[0].item()) - - return vall + # restore outer shape + self.shape = s + return time_to_freq(val, n_fft, hop_length) - def visitSifftFunc(self, ctx): - s = self.shape +def visitSifftFunc(self, ctx): + s = self.shape + hop_length = 256 + n_fft = 512 + print ("SIFFT input shape:", s) + if len(s) < 4: + raise ValueError("Input tensor must have shape (B, C, F, T_spec) for SIFFT") - hop_length = 256 - n_fft = 512 - self.shape = (s[0], - s[1], - n_fft // 2 + 1, # Frequency bins - s[2]//hop_length+1) + # Recover original time length + spec_T = s[3] + time_len = (spec_T - 1) * hop_length - shp = torch.zeros(self.shape) + # Create freq-domain variables for the subtree + shp = torch.zeros(s) + self.variables['F'] = getIndexTensorAlongDim(shp, 2) + self.variables['K'] = torch.full(s, s[2]) + self.variables['T'] = getIndexTensorAlongDim(shp, 3) + self.variables['B'] = getIndexTensorAlongDim(shp, 0) + self.variables['C'] = getIndexTensorAlongDim(shp, 1) + self.variables['R'] = torch.full_like(shp, self.variables['R'].flatten()[0].item()) - self.variables['F'] = getIndexTensorAlongDim(shp,2) - self.variables['K'] = torch.full(self.shape,self.shape[2]); + # evaluate inside ifft + val = self.visit(ctx.expr()) - self.variables['T'] = torch.full(self.shape,self.shape[3]); - self.variables['B'] = getIndexTensorAlongDim(shp, 0) - self.variables['C'] = getIndexTensorAlongDim(shp, 1) - self.variables['S'] = getIndexTensorAlongDim(shp, 3) - self.variables['R'] = torch.full_like(shp,self.variables['R'].flatten()[0].item()) + # convert back to time domain + wav = freq_to_time(val, n_fft, hop_length, time=time_len) + print("SIFFT result shape:", wav.shape) + # cleanup: remove freq-domain vars (not valid outside) + for v in ['F', 'K']: + self.variables.pop(v, None) - val = self.visit(ctx.expr()) + # restore outer shape + self.shape = (s[0], s[1], time_len) - time = s[2]; - - return freq_to_time(val, n_fft, hop_length, time) - self.shape = s; - - - self.variables['T'] = torch.full(self.shape,self.shape[2]); - self.variables['B'] = getIndexTensorAlongDim(shp, 0) - self.variables['C'] = getIndexTensorAlongDim(shp, 1) - self.variables['S'] = getIndexTensorAlongDim(shp, 2) - self.variables['R'] = torch.full_like(shp, self.variables['R'].flatten()[0].item()) - - self.variables.pop('F',None) - self.variables.pop('K',None) + return wav # Two-argument functions