numo-pocketfft 0.1.0

data/ext/numo/pocketfft/pocketfft/pocketfft.h

@@ -0,0 +1,34 @@
+/*
+ * This file is part of pocketfft.
+ * Licensed under a 3-clause BSD style license - see LICENSE.md
+ */
+
+/*! \file pocketfft.h
+ *  Public interface of the pocketfft library
+ *
+ *  Copyright (C) 2008-2018 Max-Planck-Society
+ *  \author Martin Reinecke
+ */
+
+#ifndef POCKETFFT_H
+#define POCKETFFT_H
+
+#include <stdlib.h>
+
+struct cfft_plan_i;
+typedef struct cfft_plan_i * cfft_plan;
+cfft_plan make_cfft_plan (size_t length);
+void destroy_cfft_plan (cfft_plan plan);
+int cfft_backward(cfft_plan plan, double c[], double fct);
+int cfft_forward(cfft_plan plan, double c[], double fct);
+size_t cfft_length(cfft_plan plan);
+
+struct rfft_plan_i;
+typedef struct rfft_plan_i * rfft_plan;
+rfft_plan make_rfft_plan (size_t length);
+void destroy_rfft_plan (rfft_plan plan);
+int rfft_backward(rfft_plan plan, double c[], double fct);
+int rfft_forward(rfft_plan plan, double c[], double fct);
+size_t rfft_length(rfft_plan plan);
+
+#endif

data/ext/numo/pocketfft/pocketfftext.c

@@ -0,0 +1,214 @@
+#include "pocketfftext.h"
+
+VALUE mNumo;
+VALUE mPocketfft;
+
+VALUE numo_pocketfft_fft(VALUE x_val, int is_forward)
+{
+  narray_t* x_nary;
+  double* x_pt;
+  size_t length;
+  int n_dims;
+  int n_repeats;
+  int i;
+  int res;
+  double fct;
+  VALUE z_val;
+  double* z_pt;
+  narray_t* z_nary;
+  cfft_plan plan = NULL;
+
+  if (CLASS_OF(x_val) != numo_cDComplex) {
+    x_val = rb_funcall(numo_cDComplex, rb_intern("cast"), 1, x_val);
+  }
+  if (!RTEST(nary_check_contiguous(x_val))) {
+    x_val = nary_dup(x_val);
+  }
+
+  GetNArray(x_val, x_nary);
+  n_dims = NA_NDIM(x_nary);
+  length = NA_SHAPE(x_nary)[n_dims - 1];
+  x_pt = (double*)na_get_pointer_for_read(x_val);
+
+  z_val = nary_s_new_like(numo_cDComplex, x_val);
+  z_pt = (double*)na_get_pointer_for_write(z_val);
+  GetNArray(z_val, z_nary);
+  for (i = 0; i < (int)(NA_SIZE(z_nary) * 2); z_pt[i++] = 0.0);
+
+  fct = is_forward == 1 ? 1.0 : 1.0 / length;
+  plan = make_cfft_plan(length);
+  if (!plan) {
+    return Qnil;
+  }
+
+  n_repeats = (int)(NA_SIZE(x_nary)) / length;
+  for (i = 0; i < n_repeats; i++) {
+    memcpy(z_pt, x_pt, 2 * length * sizeof(double));
+    res = is_forward == 1 ? cfft_forward(plan, z_pt, fct) : cfft_backward(plan, z_pt, fct);
+    if (res != 0) { break; }
+    z_pt += length * 2;
+    x_pt += length * 2;
+  }
+
+  if (plan) {
+    destroy_cfft_plan(plan);
+  }
+
+  return z_val;
+}
+
+/**
+ * @!visibility private
+ */
+static VALUE numo_pocketfft_cfft(VALUE self, VALUE x_val)
+{
+  return numo_pocketfft_fft(x_val, 1);
+}
+
+/**
+ * @!visibility private
+ */
+static VALUE numo_pocketfft_icfft(VALUE self, VALUE x_val)
+{
+  return numo_pocketfft_fft(x_val, 0);
+}
+
+/**
+ * @!visibility private
+ */
+static VALUE numo_pocketfft_rfft(VALUE self, VALUE x_val)
+{
+  narray_t* x_nary;
+  double* x_pt;
+  int n_dims;
+  size_t length;
+  int n_repeats;
+  int i;
+  size_t* z_shape;
+  VALUE z_val;
+  narray_t* z_nary;
+  double* z_pt;
+  int z_step;
+  rfft_plan plan = NULL;
+
+  if (CLASS_OF(x_val) != numo_cDFloat) {
+    x_val = rb_funcall(numo_cDFloat, rb_intern("cast"), 1, x_val);
+  }
+  if (!RTEST(nary_check_contiguous(x_val))) {
+    x_val = nary_dup(x_val);
+  }
+
+  GetNArray(x_val, x_nary);
+  n_dims = NA_NDIM(x_nary);
+  length = NA_SHAPE(x_nary)[n_dims - 1];
+  x_pt = (double*)na_get_pointer_for_read(x_val);
+
+  plan = make_rfft_plan(length);
+  if (!plan) {
+    return Qnil;
+  }
+
+  z_shape = ALLOCA_N(size_t, n_dims);
+  for (i = 0; i < n_dims - 1; i++) {
+    z_shape[i] = NA_SHAPE(x_nary)[i];
+  }
+  z_shape[n_dims - 1] = length / 2 + 1;
+  z_val = rb_narray_new(numo_cDComplex, n_dims, z_shape);
+  z_pt = (double*)na_get_pointer_for_write(z_val);
+  GetNArray(z_val, z_nary);
+  for (i = 0; i < (int)(NA_SIZE(z_nary) * 2); z_pt[i++] = 0.0);
+
+  z_step = (int)(NA_SHAPE(z_nary)[n_dims - 1]) * 2;
+  n_repeats = (int)(NA_SIZE(x_nary)) / length;
+  for (i = 0; i < n_repeats; i++) {
+    z_pt[z_step - 1] = 0.0;
+    memcpy(z_pt + 1, x_pt, length * sizeof(double));
+    if (rfft_forward(plan, z_pt + 1, 1.0) != 0) { break; }
+    z_pt[0] = z_pt[1];
+    z_pt[1] = 0.0;
+    z_pt += z_step;
+    x_pt += length;
+  }
+
+  if (plan) {
+    destroy_rfft_plan(plan);
+  }
+
+  return z_val;
+}
+
+/**
+ * @!visibility private
+ */
+static VALUE numo_pocketfft_irfft(VALUE self, VALUE x_val)
+{
+  narray_t* x_nary;
+  double* x_pt;
+  size_t length;
+  int n_dims;
+  int n_repeats;
+  int i;
+  double fct;
+  size_t* z_shape;
+  VALUE z_val;
+  double* z_pt;
+  narray_t* z_nary;
+  rfft_plan plan = NULL;
+
+  if (CLASS_OF(x_val) != numo_cDComplex) {
+    x_val = rb_funcall(numo_cDComplex, rb_intern("cast"), 1, x_val);
+  }
+  if (!RTEST(nary_check_contiguous(x_val))) {
+    x_val = nary_dup(x_val);
+  }
+
+  GetNArray(x_val, x_nary);
+  n_dims = NA_NDIM(x_nary);
+  length = NA_SHAPE(x_nary)[n_dims - 1];
+  x_pt = (double*)na_get_pointer_for_read(x_val);
+  fct = 1.0 / length;
+
+  plan = make_rfft_plan(length);
+  if (!plan) {
+    return Qnil;
+  }
+
+  z_shape = ALLOCA_N(size_t, n_dims);
+  for (i = 0; i < n_dims - 1; i++) {
+    z_shape[i] = NA_SHAPE(x_nary)[i];
+  }
+  z_shape[n_dims - 1] = length;
+  z_val = rb_narray_new(numo_cDFloat, n_dims, z_shape);
+  z_pt = (double*)na_get_pointer_for_write(z_val);
+  GetNArray(z_val, z_nary);
+  for (i = 0; i < (int)NA_SIZE(z_nary); z_pt[i++] = 0.0);
+
+  n_repeats = (int)(NA_SIZE(z_nary)) / length;
+  for (i = 0; i < n_repeats; i++) {
+    memcpy(z_pt + 1, x_pt + 2, (length - 1) * sizeof(double));
+    z_pt[0] = x_pt[0];
+    if (rfft_backward(plan, z_pt, fct) != 0) { break; }
+    z_pt += length;
+    x_pt += length * 2;
+  }
+
+  if (plan) {
+    destroy_rfft_plan(plan);
+  }
+
+  return z_val;
+}
+
+void Init_pocketfftext()
+{
+  rb_require("numo/narray");
+
+  mNumo = rb_define_module("Numo");
+
+  mPocketfft = rb_define_module_under(mNumo, "Pocketfft");
+
+  rb_define_module_function(mPocketfft, "ext_rfft", numo_pocketfft_rfft, 1);
+  rb_define_module_function(mPocketfft, "ext_irfft", numo_pocketfft_irfft, 1);
+  rb_define_module_function(mPocketfft, "ext_cfft", numo_pocketfft_cfft, 1);
+  rb_define_module_function(mPocketfft, "ext_icfft", numo_pocketfft_icfft, 1);
+}

data/ext/numo/pocketfft/pocketfftext.h

@@ -0,0 +1,13 @@
+#ifndef NUMO_POCKETFFT_H
+#define NUMO_POCKETFFT_H 1
+
+#include <stdlib.h>
+#include <string.h>
+
+#include <ruby.h>
+#include <numo/narray.h>
+#include <numo/template.h>
+
+#include "pocketfft/pocketfft.h"
+
+#endif /* NUMO_POCKETFFT_H */

data/lib/numo/pocketfft.rb

@@ -0,0 +1,149 @@
+# frozen_string_literal: true
+
+require 'numo/narray'
+require 'numo/pocketfft/version'
+require 'numo/pocketfft/pocketfftext'
+
+module Numo
+  module Pocketfft
+    module_function
+
+    # Compute the 1-dimensional discrete Fourier Transform.
+    # @param a [Numo::DFloat/Numo::DComplex] Real or complex 1-dimensional input array.
+    # @return [Numo::DComplex] Transformed data.
+    def fft(a)
+      raise ArgumentError, 'Expect input array to be one-dimensional.' unless a.ndim == 1
+      raw_fft(a, 0, inverse: false, real: false)
+    end
+
+    # Compute the 1-dimensional inverse discrete Fourier Transform.
+    # @param a [Numo::DComplex] Complex 1-dimensional input array.
+    # @return [Numo::DComplex] Inversed transformed data.
+    def ifft(a)
+      raise ArgumentError, 'Expect input array to be one-dimensional.' unless a.ndim == 1
+      raw_fft(a, 0, inverse: true, real: false)
+    end
+
+    # Compute the 2-dimensional discrete Fourier Transform.
+    # @param a [Numo::DFloat/Numo::DComplex] Real or complex 2-dimensional input array.
+    # @return [Numo::DComplex] Transformed data.
+    def fft2(a)
+      raise ArgumentError, 'Expect input array to be two-dimensional.' unless a.ndim == 2
+      fftn(a)
+    end
+
+    # Compute the 2-dimensional inverse discrete Fourier Transform.
+    # @param a [Numo::DComplex] Complex 2-dimensional input array.
+    # @return [Numo::DComplex] Inversed transformed data.
+    def ifft2(a)
+      raise ArgumentError, 'Expect input array to be two-dimensional.' unless a.ndim == 2
+      ifftn(a)
+    end
+
+    # Compute the N-dimensional discrete Fourier Transform.
+    # @param a [Numo::DFloat/Numo::DComplex] Real or complex input array with any-dimension.
+    # @return [Numo::DComplex] Transformed data.
+    def fftn(a)
+      b = a.dup
+      (0...b.ndim).to_a.reverse.each { |ax_id| b = raw_fft(b, ax_id, inverse: false, real: false) }
+      b
+    end
+
+    # Compute the N-dimensional inverse discrete Fourier Transform.
+    # @param a [Numo::DComplex] Complex input array with any-dimension.
+    # @return [Numo::DComplex] Inversed transformed data.
+    def ifftn(a)
+      b = a.dup
+      (0...b.ndim).to_a.each { |ax_id| b = raw_fft(b, ax_id, inverse: true, real: false) }
+      b
+    end
+
+    # Compute the 1-dimensional discrete Fourier Transform for real input.
+    # @param a [Numo::DFloat] Real 1-dimensional input array.
+    # @return [Numo::DComplex] Transformed data.
+    def rfft(a)
+      raise ArgumentError, 'Expect input array to be one-dimensional.' unless a.ndim == 1
+      raw_fft(a, 0, inverse: false, real: true)
+    end
+
+    # Compute the inverse of the 1-dimensional discrete Fourier Transform of real input.
+    # @param a [Numo::DComplex] Complex 1-dimensional input array.
+    # @return [Numo::DFloat] Inverse transformed data.
+    def irfft(a)
+      raise ArgumentError, 'Expect input array to be one-dimensional.' unless a.ndim == 1
+      raw_fft(a, 0, inverse: true, real: true)
+    end
+
+    # Compute the 2-dimensional discrete Fourier Transform for real input.
+    # @param a [Numo::DFloat] Real 2-dimensional input array.
+    # @return [Numo::DComplex] Transformed data.
+    def rfft2(a)
+      raise ArgumentError, 'Expect input array to be two-dimensional.' unless a.ndim == 2
+      rfftn(a)
+    end
+
+    # Compute the inverse of the 2-dimensional discrete Fourier Transform of real input.
+    # @param a [Numo::DComplex] Complex 2-dimensional input array.
+    # @return [Numo::DFloat] Inverse transformed data.
+    def irfft2(a)
+      raise ArgumentError, 'Expect input array to be two-dimensional.' unless a.ndim == 2
+      irfftn(a)
+    end
+
+    # Compute the N-dimensional discrete Fourier Transform for real input.
+    # @param a [Numo::DFloat] Real input array with any-dimension.
+    # @return [Numo::DComplex] Transformed data.
+    def rfftn(a)
+      last_axis_id = a.ndim - 1
+      b = raw_fft(a, last_axis_id, inverse: false, real: true)
+      (0...last_axis_id).to_a.reverse.each { |ax_id| b = raw_fft(b, ax_id, inverse: false, real: false) }
+      b
+    end
+
+    # Compute the inverse of the N-dimensional discrete Fourier Transform of real input.
+    # @param a [Numo::DComplex] Complex input array with any-dimension.
+    # @return [Numo::DFloat] Inverse transformed data.
+    def irfftn(a)
+      last_axis_id = a.ndim - 1
+      b = a.dup
+      (0...last_axis_id).to_a.each { |ax_id| b = raw_fft(b, ax_id, inverse: true, real: false) }
+      raw_fft(b, last_axis_id, inverse: true, real: true)
+    end
+
+    # @!visibility private
+    def raw_fft(a, axis_id, inverse:, real:)
+      if axis_id == a.ndim - 1
+        if real
+          if inverse
+            # zero padding
+            n = (a.shape[-1] - 1) * 2
+            b_shape = a.shape
+            b_shape[-1] = n
+            b = Numo::DComplex.zeros(*b_shape)
+            b_range = [true] * b.ndim
+            b_range[-1] = 0...a.shape[-1]
+            b[*b_range] = a
+            # inverse of dft for real data
+            ext_irfft(b)
+          else
+            ext_rfft(a)
+          end
+        else
+          if inverse
+            ext_icfft(a)
+          else
+            ext_cfft(a)
+          end
+        end
+      else
+        if inverse
+          ext_icfft(a.swapaxes(axis_id, -1)).swapaxes(axis_id, -1).dup
+        else
+          ext_cfft(a.swapaxes(axis_id, -1)).swapaxes(axis_id, -1).dup
+        end
+      end
+    end
+
+    private_class_method :ext_cfft, :ext_icfft, :ext_rfft, :ext_irfft, :raw_fft
+  end
+end

data/lib/numo/pocketfft/version.rb

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+# Numo is the top level namespace of NUmerical MOdules for Ruby.
+module Numo
+  # Numo::Pocketfft is the module that has functions for Fourier transform.
+  module Pocketfft
+    # The version of Numo::Pocketfft you are using.
+    VERSION = '0.1.0'
+  end
+end

data/numo-pocketfft.gemspec

@@ -0,0 +1,47 @@
+lib = File.expand_path('lib', __dir__)
+$LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
+require 'numo/pocketfft/version'
+
+Gem::Specification.new do |spec|
+  spec.name          = 'numo-pocketfft'
+  spec.version       = Numo::Pocketfft::VERSION
+  spec.authors       = ['yoshoku']
+  spec.email         = ['yoshoku@outlook.com']
+
+  spec.summary       = <<~MSG
+    Numo::Pocketfft provides functions for descrete Fourier Transform based on pocketfft.
+  MSG
+  spec.description = <<~MSG
+    Numo::Pocketfft provides functions for descrete Fourier Transform based on pocketfft.
+  MSG
+  spec.homepage      = 'https://github.com/yoshoku/numo-pocketfft'
+  spec.license       = 'BSD-3-Clause'
+
+  # Specify which files should be added to the gem when it is released.
+  # The `git ls-files -z` loads the files in the RubyGem that have been added into git.
+  spec.files         = Dir.chdir(File.expand_path(__dir__)) do
+    `git ls-files -z`.split("\x0").reject { |f| f.match(%r{^(test|spec|features)/}) }
+  end
+
+  # Add files in submodule: https://gist.github.com/mattconnolly/5875987
+  gem_dir = __dir__ + '/'
+  `git submodule --quiet foreach pwd`.split($OUTPUT_RECORD_SEPARATOR).each do |submodule_path|
+    Dir.chdir(submodule_path) do
+      submodule_relative_path = submodule_path.sub gem_dir, ''
+      `git ls-files`.split($OUTPUT_RECORD_SEPARATOR).each do |filename|
+        spec.files << "#{submodule_relative_path}/#{filename}"
+      end
+    end
+  end
+
+  spec.require_paths = ['lib']
+  spec.extensions    = ['ext/numo/pocketfft/extconf.rb']
+
+  spec.add_runtime_dependency 'numo-narray', '~> 0.9.1'
+
+  spec.add_development_dependency 'bundler', '~> 2.0'
+  spec.add_development_dependency 'coveralls', '~> 0.8'
+  spec.add_development_dependency 'rake', '~> 10.0'
+  spec.add_development_dependency 'rake-compiler', '~> 1.0'
+  spec.add_development_dependency 'rspec', '~> 3.0'
+end