diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml index a4a089b..771059d 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.yml +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -34,7 +34,7 @@ body: 2. Run the following OCaml code: ```ocaml let () = - let audio, _ = Soundml.Io.read Rune.c Float32 filename in + let audio, _ = Soundml.Io.read Nx.float32 filename in ignore audio ``` 3. Observe the error output. @@ -45,7 +45,7 @@ body: attributes: label: "Expected Behavior" description: "What did you expect to happen when you followed the steps above?" - placeholder: "The audio file should load into a Rune tensor without raising an exception." + placeholder: "The audio file should load into an Nx tensor without raising an exception." validations: required: true - type: textarea @@ -68,7 +68,7 @@ body: description: "Please provide details about the environment where you encountered the bug." placeholder: | - SoundML Version: 0.0.1 (or commit hash) - - OCaml Version: 5.3.0 + - OCaml Version: 5.5.0 - Operating System: Ubuntu 22.04 - type: textarea id: additional-context diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index db48df7..1fa494e 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -16,7 +16,7 @@ jobs: matrix: os: [ubuntu-latest] # , macos-latest, windows-latest] ocaml-compiler: - - "5.3.0" + - "5.5.0" runs-on: ${{ matrix.os }} @@ -38,11 +38,7 @@ jobs: - name: Install libsndfile, libsamplerate and librubberband run: | - sudo apt-get install libsndfile-dev libsamplerate0-dev librubberband-dev - - - name: Install llvm # required for rune - run: | - sudo apt-get install libllvm-20-ocaml-dev libllvm20 llvm-20 llvm-20-dev llvm-20-doc llvm-20-examples llvm-20-runtime + sudo apt-get install -y libsndfile-dev libsamplerate0-dev librubberband-dev - name: Lock the dependencies run: opam exec -- dune pkg lock diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index e446594..1cb0ecd 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -18,7 +18,7 @@ jobs: - name: Set up OCaml for linting uses: ocaml/setup-ocaml@v3 with: - ocaml-compiler: "5.3" + ocaml-compiler: "5.5" dune-cache: true - name: Run OCaml Lint & Format Check uses: ocaml/setup-ocaml/lint-fmt@v2 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 6b61abf..2637921 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -17,11 +17,7 @@ jobs: matrix: os: [ubuntu-latest] # , macos-latest, windows-latest] ocaml-compiler: - - "5.3.0" - include: - - os: ubuntu-latest - ocaml-compiler: "5.3.0" - send-coverage: true + - "5.5.0" runs-on: ${{ matrix.os }} @@ -50,11 +46,7 @@ jobs: - name: Install libsndfile, libsamplerate and librubberband run: | - sudo apt-get install libsndfile-dev libsamplerate0-dev librubberband-dev - - - name: Install llvm # required for rune - run: | - sudo apt-get install libllvm-20-ocaml-dev libllvm20 llvm-20 llvm-20-dev llvm-20-doc llvm-20-examples llvm-20-runtime + sudo apt-get install -y libsndfile-dev libsamplerate0-dev librubberband-dev - name: Install FFmpeg CLI (to generate test data) run: sudo apt-get install -y ffmpeg @@ -65,17 +57,5 @@ jobs: shell: bash - name: Run tests - if: matrix.send-coverage != true - run: | - opam exec -- dune runtest - shell: bash - - - name: Run tests (and send converage to coveralls) - if: matrix.send-coverage == true - run: | - opam exec -- dune runtest --instrument-with bisect_ppx --force - opam exec -- bisect-ppx-report send-to Coveralls + run: opam exec -- dune runtest shell: bash - env: - COVERALLS_REPO_TOKEN: ${{ secrets.GITHUB_TOKEN }} - PULL_REQUEST_NUMBER: ${{ github.event.number }} diff --git a/.ocamlformat b/.ocamlformat index cea5554..f9d42ab 100644 --- a/.ocamlformat +++ b/.ocamlformat @@ -1,3 +1,3 @@ profile = ocamlformat wrap-comments = true -version = 0.27.0 \ No newline at end of file +version = 0.28.1 \ No newline at end of file diff --git a/README.md b/README.md index ba60a51..8a5cb32 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,6 @@ [![Build](https://github.com/gabyfle/SoundML/actions/workflows/build.yml/badge.svg)](https://github.com/gabyfle/SoundML/actions/workflows/build.yml) [![Test](https://github.com/gabyfle/SoundML/actions/workflows/test.yml/badge.svg)](https://github.com/gabyfle/SoundML/actions/workflows/test.yml) -[![Coverage Status](https://coveralls.io/repos/github/gabyfle/SoundML/badge.svg)](https://coveralls.io/github/gabyfle/SoundML) ## About the Project diff --git a/doc/index.mld b/doc/index.mld index 480ada7..38e8b54 100644 --- a/doc/index.mld +++ b/doc/index.mld @@ -7,7 +7,7 @@ %} {!Soundml} is an OCaml library that enables audio processing. It provides tools to read, write, extract features and manipulate audio as time series data. -The library is built on top of the Owl library and provides a high-level API to work with audio data. It's still in an experimental stage. +The library is built on top of Raven's Nx package and provides a high-level API to work with audio data. It's still in an experimental stage. {1 API Documentation} diff --git a/dune-project b/dune-project index a1791d8..03d9cfe 100644 --- a/dune-project +++ b/dune-project @@ -16,34 +16,28 @@ (documentation https://soundml.dev) (pin - (url "git+https://github.com/raven-ml/raven.git") + (url + "git+https://github.com/raven-ml/raven.git#d9facbc227ddc80d644bdeff0d94cf76d6a0c07b") (package (name nx))) -(pin - (url "git+https://github.com/raven-ml/raven.git") - (package (name rune))) - (package (name soundml) (synopsis "An OCaml library to embed sound processing in your applications") (description - "SoundML is a library built on top of Rune to analyse sounds files. It can read, write audio, extract various features from audio files and much more.") + "SoundML is a library built on top of Raven's Nx package to analyse sound files. It can read and write audio, extract features from audio files, and much more.") (depends (ocaml - (>= 5.3.0)) + (>= 5.5.0)) dune (conf-sndfile :build) (conf-rubberband :build) ; Samplerate is already a requirement of rubberband (conf-samplerate :build) (dune-configurator :build) - rune + (nx + (= dev)) (odoc :with-doc) (alcotest :with-test) - (yojson :with-test) - (bisect_ppx - (and - :with-test - (>= 2.5.0)))) + (yojson :with-test)) (tags (sound spectrogram fourier digital-filters))) diff --git a/soundml.opam b/soundml.opam index 9af38c7..6b0e145 100644 --- a/soundml.opam +++ b/soundml.opam @@ -2,7 +2,7 @@ opam-version: "2.0" synopsis: "An OCaml library to embed sound processing in your applications" description: - "SoundML is a library built on top of Rune to analyse sounds files. It can read, write audio, extract various features from audio files and much more." + "SoundML is a library built on top of Raven's Nx package to analyse sound files. It can read and write audio, extract features from audio files, and much more." maintainer: ["Gabriel Santamaria "] authors: ["Gabriel Santamaria "] license: "Apache-2.0" @@ -11,17 +11,16 @@ homepage: "https://github.com/gabyfle/SoundML" doc: "https://soundml.dev" bug-reports: "https://github.com/gabyfle/SoundML/issues" depends: [ - "ocaml" {>= "5.3.0"} + "ocaml" {>= "5.5.0"} "dune" {>= "3.17"} "conf-sndfile" {build} "conf-rubberband" {build} "conf-samplerate" {build} "dune-configurator" {build} - "rune" + "nx" {= "dev"} "odoc" {with-doc} "alcotest" {with-test} "yojson" {with-test} - "bisect_ppx" {with-test & >= "2.5.0"} ] build: [ ["dune" "subst"] {dev} @@ -38,3 +37,9 @@ build: [ ] ] dev-repo: "git+https://github.com/gabyfle/SoundML.git" +pin-depends: [ + [ + "nx.dev" + "git+https://github.com/raven-ml/raven.git#d9facbc227ddc80d644bdeff0d94cf76d6a0c07b" + ] +] diff --git a/soundml.opam.template b/soundml.opam.template new file mode 100644 index 0000000..998218c --- /dev/null +++ b/soundml.opam.template @@ -0,0 +1,6 @@ +pin-depends: [ + [ + "nx.dev" + "git+https://github.com/raven-ml/raven.git#d9facbc227ddc80d644bdeff0d94cf76d6a0c07b" + ] +] diff --git a/src/dune b/src/dune index b2a881b..0c56f63 100644 --- a/src/dune +++ b/src/dune @@ -2,50 +2,38 @@ (name aformat) (package soundml) (modules aformat) - (instrumentation - (backend bisect_ppx)) (wrapped false)) (library (name window) (package soundml) - (libraries rune utils) + (libraries nx utils) (modules window) - (instrumentation - (backend bisect_ppx)) (wrapped false)) (library (name utils) (package soundml) - (libraries rune) + (libraries nx) (modules utils) - (instrumentation - (backend bisect_ppx)) (wrapped false)) (library (name transform) (package soundml) (modules transform) - (libraries rune window utils) - (instrumentation - (backend bisect_ppx)) + (libraries nx window utils) (wrapped false)) (library (name filterbank) (package soundml) (modules filterbank) - (libraries rune utils) - (instrumentation - (backend bisect_ppx)) + (libraries nx utils) (wrapped false)) (library (name soundml) (public_name soundml) (modules soundml) - (libraries aformat window io utils transform filterbank) - (instrumentation - (backend bisect_ppx))) + (libraries aformat window io utils transform filterbank)) diff --git a/src/effects/dune b/src/effects/dune index d381447..ae7c0ae 100644 --- a/src/effects/dune +++ b/src/effects/dune @@ -1,8 +1,6 @@ ;(library ; (name effects) ; (public_name soundml.effects) -; (instrumentation -; (backend bisect_ppx)) ; (modules ; time ; filter diff --git a/src/feature/dune b/src/feature/dune index 2d2025c..03b810b 100644 --- a/src/feature/dune +++ b/src/feature/dune @@ -1,7 +1,5 @@ ;(library ; (name feature) ; (public_name soundml.feature) -; (instrumentation -; (backend bisect_ppx)) ; (libraries audio utils window nx) ; (modules spectral)) diff --git a/src/filterbank.ml b/src/filterbank.ml index 64efa9c..670bb32 100644 --- a/src/filterbank.ml +++ b/src/filterbank.ml @@ -23,14 +23,13 @@ type normalization = Slaney | Power_norm of float let mel_filterbank ?f_max ?(htk : bool = false) ?norm ~sample_rate ~n_fft - ~n_mels ~f_min (device : 'dev Rune.device) (dtype : (float, 'b) Rune.dtype) - = + ~n_mels ~f_min (dtype : (float, 'b) Nx.dtype) = (* Input validation *) if sample_rate <= 0 then invalid_arg "sample_rate must be positive" ; if n_fft <= 0 then invalid_arg "n_fft must be positive" ; if n_mels <= 0 then invalid_arg "n_mels must be positive" ; if f_min < 0.0 then invalid_arg "f_min must be non-negative" ; - if n_mels = 0 then Rune.empty device dtype [|0; (n_fft / 2) + 1|] + if n_mels = 0 then Nx.empty dtype [|0; (n_fft / 2) + 1|] else let f_max = match f_max with @@ -48,52 +47,46 @@ let mel_filterbank ?f_max ?(htk : bool = false) ?norm ~sample_rate ~n_fft let freq_step = 1.0 /. (float_of_int n_fft *. d) in let fftfreqs = let indices = Array.init n_freqs (fun i -> float_of_int i *. freq_step) in - Rune.create device dtype [|n_freqs|] indices + Nx.create dtype [|n_freqs|] indices in let mel_freqs = - Utils.melfreqs ~n_mels:(n_mels + 2) ~f_min ~f_max ~htk device dtype + Utils.melfreqs ~n_mels:(n_mels + 2) ~f_min ~f_max ~htk dtype in let fdiff = - let n = Rune.size mel_freqs in - Rune.sub - (Rune.slice [R (1, n)] mel_freqs) - (Rune.slice [R (0, n - 1)] mel_freqs) + let n = Nx.numel mel_freqs in + Nx.sub (Nx.slice [R (1, n)] mel_freqs) (Nx.slice [R (0, n - 1)] mel_freqs) in - let ramps = Utils.outer Rune.sub mel_freqs fftfreqs in + let ramps = Utils.outer Nx.sub mel_freqs fftfreqs in let lower = - Rune.div - (Rune.neg (Rune.slice [R (0, n_mels)] ramps)) - (Rune.reshape [|n_mels; 1|] (Rune.slice [R (0, n_mels)] fdiff)) + Nx.div + (Nx.neg (Nx.slice [R (0, n_mels)] ramps)) + (Nx.reshape [|n_mels; 1|] (Nx.slice [R (0, n_mels)] fdiff)) in let upper = - Rune.div - (Rune.slice [R (2, n_mels + 2)] ramps) - (Rune.reshape [|n_mels; 1|] (Rune.slice [R (1, n_mels + 1)] fdiff)) + Nx.div + (Nx.slice [R (2, n_mels + 2)] ramps) + (Nx.reshape [|n_mels; 1|] (Nx.slice [R (1, n_mels + 1)] fdiff)) in (* Intersect slopes *) - let weights = - Rune.maximum (Rune.zeros_like lower) (Rune.minimum lower upper) - in + let weights = Nx.maximum (Nx.zeros_like lower) (Nx.minimum lower upper) in let weights = match norm with | Some Slaney -> let enorm = - Rune.div - (Rune.scalar device dtype 2.0) - (Rune.sub - (Rune.slice [R (2, n_mels + 2)] mel_freqs) - (Rune.slice [R (0, n_mels)] mel_freqs) ) + Nx.div (Nx.scalar dtype 2.0) + (Nx.sub + (Nx.slice [R (2, n_mels + 2)] mel_freqs) + (Nx.slice [R (0, n_mels)] mel_freqs) ) in - let enorm = Rune.reshape [|n_mels; 1|] enorm in - Rune.mul weights enorm + let enorm = Nx.reshape [|n_mels; 1|] enorm in + Nx.mul weights enorm | Some (Power_norm p) -> let norm = - Rune.pow_s - (Rune.sum ~axes:[|-1|] ~keepdims:true - (Rune.pow_s (Rune.abs weights) p) ) + Nx.pow_s + (Nx.sum ~axes:[-1] ~keepdims:true (Nx.pow_s (Nx.abs weights) p)) (1. /. p) in - Rune.div weights (Rune.add_s norm 1e-8) + Nx.div weights (Nx.add_s norm 1e-8) | None -> weights in diff --git a/src/filterbank.mli b/src/filterbank.mli index c28b486..e8d9866 100644 --- a/src/filterbank.mli +++ b/src/filterbank.mli @@ -43,10 +43,9 @@ val mel_filterbank : -> n_fft:int -> n_mels:int -> f_min:float - -> 'dev Rune.device - -> (float, 'b) Rune.dtype - -> (float, 'b, 'dev) Rune.t -(** [mel_filterbank dtype ~sample_rate ~n_fft ~n_mels ~f_min] creates a mel-scale filterbank. + -> (float, 'b) Nx.dtype + -> (float, 'b) Nx.t +(** [mel_filterbank ~sample_rate ~n_fft ~n_mels ~f_min dtype] creates a mel-scale filterbank. A mel filterbank is a collection of triangular filters spaced according to the mel scale, which approximates human auditory perception. Each filter has a @@ -60,8 +59,7 @@ val mel_filterbank : @param n_fft FFT size, determines frequency resolution (must be positive) @param n_mels Number of mel bands to generate (must be positive) @param f_min Minimum frequency in Hz (must be non-negative and < f_max) - @param device Rune device to use for computation (default: automatic selection) - @param dtype Data type for the filterbank (e.g., Rune.float32, Rune.float64) + @param dtype Data type for the filterbank (e.g., Nx.float32, Nx.float64) @return Mel filterbank matrix of shape [n_mels; n_fft/2 + 1] @raise Invalid_argument if sample_rate <= 0 @@ -75,17 +73,16 @@ val mel_filterbank : Basic usage with automatic device selection: {[ let filterbank = Filterbank.mel_filterbank - Rune.float32 ~sample_rate:22050 ~n_fft:1024 ~n_mels:128 - ~f_min:0.0 in + ~f_min:0.0 + Nx.float32 in (* filterbank has shape [128; 513] *) ]} - With explicit device and normalization: + With Slaney normalization: {[ - let device = Rune.metal () in let filterbank = Filterbank.mel_filterbank ~norm:Slaney ~f_max:8000.0 @@ -93,8 +90,7 @@ val mel_filterbank : ~n_fft:512 ~n_mels:80 ~f_min:80.0 - device - Rune.float64 + Nx.float64 in (* filterbank has shape [80; 257] with Slaney normalization *) ]} @@ -107,8 +103,7 @@ val mel_filterbank : ~n_fft:2048 ~n_mels:40 ~f_min:20.0 - Rune.float32 - Rune.ocaml + Nx.float32 in (* filterbank uses HTK mel scale formula *) ]} @@ -126,7 +121,5 @@ val mel_filterbank : {3 Performance Notes} - - Use Metal device for best performance on Apple Silicon - - Use C device for good performance on other platforms - Float32 is usually sufficient precision and faster than Float64 - Larger n_fft values provide better frequency resolution but are slower *) diff --git a/src/io/cio/read.hxx b/src/io/cio/read.hxx index 2dc250f..7a0f45f 100644 --- a/src/io/cio/read.hxx +++ b/src/io/cio/read.hxx @@ -260,7 +260,7 @@ namespace SoundML template inline value caml_read_audio_file(value filename, value res_typ, value trgt_sr) { - CAMLparam0(); + CAMLparam3(filename, res_typ, trgt_sr); CAMLlocal3(audio_array, caml_sample_rate, returns); using namespace SoundML::IO; @@ -290,7 +290,7 @@ inline value caml_read_audio_file(value filename, value res_typ, value trgt_sr) } else { - trgt_sr = sample_rate; + trgt_sr_val = sample_rate; reader = new SndfileReader(nframes, channels, sample_rate); } @@ -342,7 +342,7 @@ inline value caml_read_audio_file(value filename, value res_typ, value trgt_sr) Store_field(returns, 0, audio_array); Store_field(returns, 1, caml_sample_rate); - return returns; + CAMLreturn(returns); } #endif /* SOUNDFILE_READER_H */ diff --git a/src/io/dune b/src/io/dune index 185b0a3..7946ce8 100644 --- a/src/io/dune +++ b/src/io/dune @@ -15,9 +15,7 @@ (c_library_flags (:include c_library_flags.sexp)) (modules io) - (libraries aformat utils) - (instrumentation - (backend bisect_ppx)) + (libraries aformat nx utils) (wrapped true)) (rule diff --git a/src/io/io.ml b/src/io/io.ml index e7238e5..def3823 100644 --- a/src/io/io.ml +++ b/src/io/io.ml @@ -59,13 +59,13 @@ external caml_read_audio_file_f64 : -> (float, Bigarray.float64_elt, Bigarray.c_layout) Bigarray.Genarray.t * int = "caml_read_audio_file_f64" -let to_mono (audio_tensor : (float, 'a, 'dev) Rune.t) = +let to_mono (audio_tensor : (float, 'a) Nx.t) = try - if Rune.ndim audio_tensor > 1 then - let shape = Rune.shape audio_tensor in + if Nx.ndim audio_tensor > 1 then + let shape = Nx.shape audio_tensor in if Array.length shape < 2 then raise (Invalid_argument "Invalid tensor shape for mono conversion") - else Rune.mean ~axes:[|1|] ~keepdims:false audio_tensor + else Nx.mean ~axes:[1] ~keepdims:false audio_tensor else audio_tensor with | Invalid_argument _ as e -> @@ -79,28 +79,26 @@ let read : type a. ?res_typ:resampling_t -> ?sample_rate:int -> ?mono:bool - -> 'dev Rune.device - -> (float, a) Rune.dtype + -> (float, a) Nx.dtype -> string - -> (float, a, 'dev) Rune.t * int = + -> (float, a) Nx.t * int = fun ?(res_typ : resampling_t = SOXR_HQ) ?(sample_rate : int = 22050) - ?(mono : bool = true) (device : 'dev Rune.device) typ - (filename : string) -> + ?(mono : bool = true) typ (filename : string) -> if sample_rate <= 0 then raise (Invalid_argument "Sample rate must be positive") ; if String.length filename = 0 then raise (Invalid_argument "Filename cannot be empty") ; let read_func : type a. - (float, a) Rune.dtype + (float, a) Nx.dtype -> string -> resampling_t -> int -> (float, a, Bigarray.c_layout) Bigarray.Genarray.t * int = fun typ -> match typ with - | Rune.Float32 -> + | Nx.Float32 -> caml_read_audio_file_f32 - | Rune.Float64 -> + | Nx.Float64 -> caml_read_audio_file_f64 | _ -> raise @@ -113,13 +111,13 @@ let read : type a. let data_shape = Bigarray.Genarray.dims data in if Array.length data_shape = 0 || Array.fold_left ( * ) 1 data_shape = 0 then raise (Invalid_format "Audio file contains no data") ; - let data = Rune.of_bigarray device data in + let data = Nx.of_bigarray data in (* Apply mono conversion if requested *) let data = if mono then to_mono data else data in (* Transpose if multi-dimensional (channels first to channels last) *) let data = - if Rune.ndim data > 1 then - try Rune.transpose data + if Nx.ndim data > 1 then + try Nx.transpose data with exn -> raise (Internal_error ("Transpose failed: " ^ Printexc.to_string exn)) else data @@ -152,23 +150,23 @@ external caml_write_audio_file_f64 : -> unit = "caml_write_audio_file_f64" let write : type a. - ?format:Aformat.t -> string -> (float, a, 'dev) Rune.t -> int -> unit = - fun ?format (filename : string) (x : (float, a, 'dev) Rune.t) sample_rate -> + ?format:Aformat.t -> string -> (float, a) Nx.t -> int -> unit = + fun ?format (filename : string) (x : (float, a) Nx.t) sample_rate -> if sample_rate <= 0 then raise (Invalid_argument "Sample rate must be positive") ; if String.length filename = 0 then raise (Invalid_argument "Filename cannot be empty") ; let write_func : type a. - (float, a) Rune.dtype + (float, a) Nx.dtype -> string -> (float, a, Bigarray.c_layout) Bigarray.Genarray.t -> int * int * int * int -> unit = fun typ -> match typ with - | Rune.Float32 -> + | Nx.Float32 -> caml_write_audio_file_f32 - | Rune.Float64 -> + | Nx.Float64 -> caml_write_audio_file_f64 | _ -> raise @@ -186,13 +184,13 @@ let write : type a. else Option.get format in let format = Aformat.to_int format in - let data = Rune.transpose x in - let dshape = Rune.shape data in + let data = Nx.transpose x in + let dshape = Nx.shape data in let nframes = dshape.(0) in let channels = if Array.length dshape > 1 then dshape.(1) else 1 in try - let dtype = Rune.dtype data in - write_func dtype filename (Rune.to_bigarray data) + let dtype = Nx.dtype data in + write_func dtype filename (Nx.to_bigarray data) (nframes, sample_rate, channels, format) with | ( File_not_found _ diff --git a/src/io/io.mli b/src/io/io.mli index 7fdde5f..8f95615 100644 --- a/src/io/io.mli +++ b/src/io/io.mli @@ -58,11 +58,10 @@ val read : ?res_typ:resampling_t -> ?sample_rate:int -> ?mono:bool - -> 'dev Rune.device - -> (float, 'a) Rune.dtype + -> (float, 'a) Nx.dtype -> string - -> (float, 'a, 'dev) Rune.t * int -(** [read device dtype filename] reads an audio file and returns audio data with sample rate. + -> (float, 'a) Nx.t * int +(** [read dtype filename] reads an audio file and returns audio data with sample rate. Loads audio data from various file formats with optional resampling and channel conversion. The function automatically handles format detection @@ -71,8 +70,7 @@ val read : @param resampling Resampling quality (default: High_quality) @param target_sample_rate Target sample rate in Hz (default: 22050) @param force_mono Convert to mono by averaging channels (default: true) - @param device Device on which the resulting tensor should be loaded. - @param dtype Data type for audio samples (Rune.float32 or Rune.float64) + @param dtype Data type for audio samples (Nx.float32 or Nx.float64) @param filename Path to the audio file to read @return Tuple of (audio_tensor, actual_sample_rate) @@ -84,30 +82,28 @@ val read : Basic audio loading: {[ - let audio_data, sample_rate = IO.read ~device:Rune.ocaml - Rune.float32 - ~filename:"audio.wav" in + let audio_data, sample_rate = Io.read Nx.float32 "audio.wav" in (* audio_data is mono float32 tensor at 22050 Hz *) ]} Load stereo audio without resampling: {[ - let audio_data, sample_rate = IO.read_audio_file - ~resampling:No_resampling - ~force_mono:false - Rune.float64 - ~filename:"stereo.flac" in + let audio_data, sample_rate = Io.read + ~res_typ:Io.NONE + ~mono:false + Nx.float64 + "stereo.flac" in (* Preserves original sample rate and stereo channels *) ]} - Load with specific device and sample rate: + Load with a specific sample rate: {[ - let audio_data, sample_rate = IO.read - ~target_sample_rate:44100 - ~resampling:Very_high_quality - Rune.float32 - ~filename:"music.mp3" in - (* High-quality resampling to 44.1 kHz on Metal device *) + let audio_data, sample_rate = Io.read + ~sample_rate:44100 + ~res_typ:Io.SOXR_VHQ + Nx.float32 + "music.mp3" in + (* Very-high-quality resampling to 44.1 kHz *) ]} {3 Supported Formats} @@ -127,8 +123,7 @@ val read : (** {2 Audio File Writing} *) -val write : - ?format:Aformat.t -> string -> (float, 'a, 'dev) Rune.t -> int -> unit +val write : ?format:Aformat.t -> string -> (float, 'a) Nx.t -> int -> unit (** [write filename audio_data sample_rate] writes audio data to a file. Saves audio tensor data to various file formats with automatic format diff --git a/src/transform.ml b/src/transform.ml index 0c56846..b7e03a8 100644 --- a/src/transform.ml +++ b/src/transform.ml @@ -19,66 +19,45 @@ (* *) (*****************************************************************************) -let stft ?(window = `Hanning) ?(win_length = 2048) ?(center = true) ~n_fft - ~hop_length (x : (float, 'a, 'dev) Rune.t) = +let stft ?(window = `Hanning) ?win_length ?(center = true) ~n_fft ~hop_length + (x : (float, 'a) Nx.t) = + let win_length = Option.value win_length ~default:n_fft in (* Input validation *) if n_fft <= 0 then invalid_arg "n_fft must be positive" ; if hop_length <= 0 then invalid_arg "hop_length must be positive" ; if win_length <= 0 then invalid_arg "win_length must be positive" ; if win_length > n_fft then invalid_arg "win_length cannot be larger than n_fft" ; - let device = Rune.device x in - let dtype = Rune.dtype x in + let dtype = Nx.dtype x in (* 1. Create window and pad it if necessary to n_fft *) - let fft_window = Window.get window device dtype ~fftbins:true win_length in + let fft_window = Window.get window dtype ~fftbins:true win_length in let fft_window = if win_length < n_fft then let pad_width = (n_fft - win_length) / 2 in - Rune.pad [|(pad_width, n_fft - win_length - pad_width)|] 0.0 fft_window + Nx.pad [|(pad_width, n_fft - win_length - pad_width)|] 0.0 fft_window else fft_window in (* 2. Pad the input signal for centering *) let x_padded = if center then ( let pad_width = n_fft / 2 in - let rank = Rune.ndim x in + let rank = Nx.ndim x in let padding = Array.make rank (0, 0) in padding.(rank - 1) <- (pad_width, pad_width) ; - Rune.pad padding 0.0 x ) + Nx.pad padding 0.0 x ) else x in - let padded_shape = Rune.shape x_padded in + let padded_shape = Nx.shape x_padded in let signal_length = padded_shape.(Array.length padded_shape - 1) in if n_fft > signal_length then invalid_arg (Printf.sprintf "n_fft=%d is too large for input signal of length=%d" n_fft signal_length ) ; - (* 3. Calculate number of frames *) - let num_frames = 1 + ((signal_length - n_fft) / hop_length) in - let frame_indices = Rune.arange device Rune.int32 0 num_frames 1 in - (* 4. Define the function to be vmapped *) - let process_single_frame frame_index = - let start = Rune.mul_s frame_index (Int32.of_int hop_length) in - let start_reshaped = Rune.unsqueeze ~axes:[|1|] start in - let offsets = Rune.arange device Rune.int32 0 n_fft 1 in - let indices = Rune.add start_reshaped offsets in - (* Dynamic slice using take. Assuming take flattens the taken dimensions. *) - let frame_flat = Rune.take ~axis:(-1) indices x_padded in - (* Reshape back to (..., num_frames, n_fft) using -1 for num_frames *) - let x_rank = Rune.ndim x_padded in - let other_dims = Array.sub (Rune.shape x_padded) 0 (x_rank - 1) in - let target_shape = Array.append other_dims [|-1; n_fft|] in - let frame = Rune.reshape target_shape frame_flat in - (* Apply window. *) - let windowed_frame = Rune.mul frame fft_window in - (* FFT on the last axis *) - Rune.rfft ~axis:(-1) windowed_frame + let frames = + Nx.extract_patches ~kernel_size:[|n_fft|] ~stride:[|hop_length|] + ~dilation:[|1|] + ~padding:[|(0, 0)|] + x_padded in - (* 5. vmap the function over the frame indices *) - let stft_matrix = Rune.vmap process_single_frame frame_indices in - (* 6. Transpose to (..., freq, time) by moving the first axis (frames) to the - end *) - let rank = Rune.ndim stft_matrix in - let axes = Array.init rank (fun i -> if i = rank - 1 then 0 else i + 1) in - let stft_transposed = Rune.transpose ~axes stft_matrix in - stft_transposed + let fft_window = Nx.reshape [|n_fft; 1|] fft_window in + Nx.rfft ~axis:(-2) (Nx.mul frames fft_window) diff --git a/src/transform.mli b/src/transform.mli index 7e9dcce..3ae4453 100644 --- a/src/transform.mli +++ b/src/transform.mli @@ -32,8 +32,8 @@ val stft : -> ?center:bool -> n_fft:int -> hop_length:int - -> (float, 'a, 'dev) Rune.t - -> (Complex.t, Bigarray_ext.complex64_elt, 'dev) Rune.t + -> (float, 'a) Nx.t + -> Nx.complex128_t (** [stft ~n_fft ~hop_length x] computes the Short-Time Fourier Transform. The STFT transforms a time-domain signal into a time-frequency representation @@ -45,13 +45,14 @@ val stft : @param center Whether to center the signal by padding (default: true) @param n_fft FFT size, determines frequency resolution (must be positive) @param hop_length Number of samples between successive frames (must be positive) - @param x Input audio signal (1D tensor) - @return Complex STFT matrix of shape [n_fft/2 + 1; n_frames] + @param x Input audio tensor whose last axis contains the samples + @return Complex STFT tensor of shape [...; n_fft/2 + 1; n_frames] @raise Invalid_argument if n_fft <= 0 @raise Invalid_argument if hop_length <= 0 + @raise Invalid_argument if win_length <= 0 @raise Invalid_argument if win_length > n_fft (when specified) - @raise Invalid_argument if signal is not 1D + @raise Invalid_argument if n_fft is larger than the padded signal {3 Examples} @@ -87,7 +88,6 @@ val stft : {3 Performance Notes} - - Use Metal device for best performance on Apple Silicon - Powers of 2 for n_fft are most efficient (e.g., 512, 1024, 2048) - Larger hop_length reduces computation but may miss transient events - Consider using float32 input for better performance *) diff --git a/src/utils.ml b/src/utils.ml index e213cf1..55c6cbd 100644 --- a/src/utils.ml +++ b/src/utils.ml @@ -21,140 +21,127 @@ module Convert = struct let mel_to_hz ?(htk = false) mels = - let dtype = Rune.dtype mels in + let dtype = Nx.dtype mels in if htk then - let term = Rune.div_s mels 2595. in - let term = Rune.pow (Rune.scalar (Rune.device mels) dtype 10.) term in - let term = Rune.sub_s term 1. in - Rune.mul_s term 700. + let term = Nx.div_s mels 2595. in + let term = Nx.pow (Nx.scalar dtype 10.) term in + let term = Nx.sub_s term 1. in + Nx.mul_s term 700. else let f_min = 0.0 in let f_sp = 200.0 /. 3. in let min_log_hz = 1000. in let min_log_mel = (min_log_hz -. f_min) /. f_sp in let logstep = Float.log 6.4 /. 27.0 in - let linear_mask = - Rune.less_equal mels (Rune.scalar (Rune.device mels) dtype min_log_mel) - in - let linear_result = Rune.(add_s (mul_s mels f_sp) f_min) in + let linear_mask = Nx.less_equal mels (Nx.scalar dtype min_log_mel) in + let linear_result = Nx.(add_s (mul_s mels f_sp) f_min) in let log_result = - Rune.(mul_s (exp (mul_s (sub_s mels min_log_mel) logstep)) min_log_hz) + Nx.(mul_s (exp (mul_s (sub_s mels min_log_mel) logstep)) min_log_hz) in - Rune.where linear_mask linear_result log_result + Nx.where linear_mask linear_result log_result let hz_to_mel ?(htk = false) freqs = - let dtype = Rune.dtype freqs in + let dtype = Nx.dtype freqs in if htk then - let term = Rune.div_s freqs 700. in - let term = Rune.add_s term 1. in - let term = Rune.log term in - let term = Rune.div_s term (Float.log 10.) in - Rune.mul_s term 2595. + let term = Nx.div_s freqs 700. in + let term = Nx.add_s term 1. in + let term = Nx.log term in + let term = Nx.div_s term (Float.log 10.) in + Nx.mul_s term 2595. else let f_min = 0.0 in let f_sp = 200.0 /. 3. in let min_log_hz = 1000. in let min_log_mel = (min_log_hz -. f_min) /. f_sp in let logstep = Float.log 6.4 /. 27.0 in - let linear_mask = - Rune.less_equal freqs (Rune.scalar (Rune.device freqs) dtype min_log_hz) - in - let linear_result = Rune.(div_s (sub_s freqs f_min) f_sp) in + let linear_mask = Nx.less_equal freqs (Nx.scalar dtype min_log_hz) in + let linear_result = Nx.(div_s (sub_s freqs f_min) f_sp) in let log_result = - Rune.(add_s (div_s (log (div_s freqs min_log_hz)) logstep) min_log_mel) + Nx.(add_s (div_s (log (div_s freqs min_log_hz)) logstep) min_log_mel) in - let res = Rune.where linear_mask linear_result log_result in + let res = Nx.where linear_mask linear_result log_result in (* Handle NaN values by replacing them with 0.0 *) - let nan_mask = Rune.isnan res in - let zero_tensor = Rune.scalar (Rune.device res) (Rune.dtype res) 0.0 in - Rune.where nan_mask zero_tensor res + let nan_mask = Nx.isnan res in + let zero_tensor = Nx.scalar (Nx.dtype res) 0.0 in + Nx.where nan_mask zero_tensor res - type ('a, 'dev) reference = + type 'a reference = | RefFloat of float - | RefFunction of ((float, 'a, 'dev) Rune.t -> float) + | RefFunction of ((float, 'a) Nx.t -> float) - let power_to_db ?(amin = 1e-10) ?(top_db = 80.) ref - (s : (float, 'a, 'dev) Rune.t) = + let power_to_db ?(amin = 1e-10) ?(top_db = 80.) ref (s : (float, 'a) Nx.t) = assert (amin > 0.) ; let ref_value = match ref with RefFloat x -> x | RefFunction f -> f s in - let log_spec = Rune.mul_s (Rune.div_s (Rune.log s) (Float.log 10.)) 10. in + let log_spec = Nx.mul_s (Nx.div_s (Nx.log s) (Float.log 10.)) 10. in let log_spec = - Rune.( + Nx.( sub log_spec (mul_s (div_s (log (maximum - (scalar (device s) (dtype s) amin) - (scalar (device s) (dtype s) ref_value) ) ) + (scalar (dtype s) amin) + (scalar (dtype s) ref_value) ) ) (Float.log 10.) ) 10. ) ) in assert (top_db >= 0.0) ; let finite_values = - Rune.where (Rune.isfinite log_spec) log_spec - (Rune.scalar (Rune.device log_spec) (Rune.dtype log_spec) (-1e8)) + Nx.where (Nx.isfinite log_spec) log_spec + (Nx.scalar (Nx.dtype log_spec) (-1e8)) in - let max_val = Rune.max finite_values in - let max_scalar = Rune.item [] max_val in - let threshold = - Rune.scalar (Rune.device log_spec) (Rune.dtype log_spec) - (max_scalar -. top_db) - in - Rune.maximum log_spec threshold + let max_val = Nx.max finite_values in + let max_scalar = Nx.item [] max_val in + let threshold = Nx.scalar (Nx.dtype log_spec) (max_scalar -. top_db) in + Nx.maximum log_spec threshold - let db_to_power ?(amin = 1e-10) (ref : ('a, 'dev) reference) - (s : (float, 'a, 'dev) Rune.t) = + let db_to_power ?(amin = 1e-10) (ref : 'a reference) (s : (float, 'a) Nx.t) = assert (amin > 0.) ; let ref_value = match ref with RefFloat x -> x | RefFunction f -> f s in - let amin_t = Rune.scalar (Rune.device s) (Rune.dtype s) amin in - let ref_value_t = Rune.scalar (Rune.device s) (Rune.dtype s) ref_value in + let amin_t = Nx.scalar (Nx.dtype s) amin in + let ref_value_t = Nx.scalar (Nx.dtype s) ref_value in let log_ref = - Rune.mul_s - (Rune.div_s - (Rune.log (Rune.maximum amin_t ref_value_t)) - (Float.log 10.) ) + Nx.mul_s + (Nx.div_s (Nx.log (Nx.maximum amin_t ref_value_t)) (Float.log 10.)) 10. in - let spec = Rune.add s log_ref in - let spec = Rune.div_s spec 10. in - Rune.pow (Rune.scalar (Rune.device s) (Rune.dtype s) 10.) spec + let spec = Nx.add s log_ref in + let spec = Nx.div_s spec 10. in + Nx.pow (Nx.scalar (Nx.dtype s) 10.) spec end let melfreqs ?(n_mels = 128) ?(f_min = 0.) ?(f_max = 11025.) ?(htk = false) - (device : 'dev Rune.device) (kd : (float, 'b) Rune.dtype) = + (kd : (float, 'b) Nx.dtype) = (* Input validation *) if n_mels <= 0 then invalid_arg "n_mels must be positive" ; if f_min < 0.0 then invalid_arg "f_min must be non-negative" ; if f_min >= f_max then invalid_arg "f_min must be less than f_max" ; - let bounds = - Rune.create device kd [|2|] [|f_min; f_max|] |> Convert.hz_to_mel ~htk - in + let bounds = Nx.create kd [|2|] [|f_min; f_max|] |> Convert.hz_to_mel ~htk in let mel_f = - Rune.linspace device kd ~endpoint:true (Rune.item [0] bounds) - (Rune.item [1] bounds) n_mels + Nx.linspace kd ~endpoint:true (Nx.item [0] bounds) (Nx.item [1] bounds) + n_mels in Convert.mel_to_hz mel_f ~htk let outer op x y = (* Generalized outer product using reshape and broadcasting *) - let x_shape = Rune.shape x in - let y_shape = Rune.shape y in + let x_shape = Nx.shape x in + let y_shape = Nx.shape y in let x_ndim = Array.length x_shape in let y_ndim = Array.length y_shape in (* Reshape x to add singleton dimensions for y's dimensions *) (* x: [x_shape] -> [x_shape, 1, 1, ..., 1] (y_ndim ones) *) let x_new_shape = Array.concat [x_shape; Array.make y_ndim 1] in - let x_reshaped = Rune.reshape x_new_shape x in + let x_reshaped = Nx.reshape x_new_shape x in (* Reshape y to add singleton dimensions for x's dimensions *) (* y: [y_shape] -> [1, 1, ..., 1, y_shape] (x_ndim ones) *) let y_new_shape = Array.concat [Array.make x_ndim 1; y_shape] in - let y_reshaped = Rune.reshape y_new_shape y in + let y_reshaped = Nx.reshape y_new_shape y in (* Apply the operation - broadcasting will handle the rest *) op x_reshaped y_reshaped let pad_center signal ~size ~pad_value = - let signal_shape = Rune.shape signal in + let signal_shape = Nx.shape signal in let signal_length = signal_shape.(0) in if size < signal_length then invalid_arg "size must be >= signal length" ; if size = signal_length then signal @@ -163,4 +150,4 @@ let pad_center signal ~size ~pad_value = let pad_left = pad_total / 2 in let pad_right = pad_total - pad_left in let padding = [|(pad_left, pad_right)|] in - Rune.pad padding pad_value signal + Nx.pad padding pad_value signal diff --git a/src/utils.mli b/src/utils.mli index c4e9195..e47ed67 100644 --- a/src/utils.mli +++ b/src/utils.mli @@ -27,15 +27,13 @@ (** {2 Signal Framing and Windowing} *) -val pad_center : - ('a, 'b, 'dev) Rune.t -> size:int -> pad_value:'a -> ('a, 'b, 'dev) Rune.t +val pad_center : ('a, 'b) Nx.t -> size:int -> pad_value:'a -> ('a, 'b) Nx.t (** [pad_center signal ~size ~pad_value] pads a signal to center it. Centers the signal by padding it symmetrically on both sides. This is commonly used in STFT preprocessing to ensure the first and last frames are properly centered on the signal boundaries. - @param device Rune device to use for computation (default: same as input) @param signal Input signal to pad @param size Total desired size after padding @param pad_value Value to use for padding @@ -47,7 +45,7 @@ val pad_center : {3 Example} {[ - let signal = Rune.ones Rune.float32 [|100|] in + let signal = Nx.ones Nx.float32 [|100|] in let padded = Utils.pad_center signal ~size:200 ~pad_value:0.0 in (* padded has 50 zeros, then 100 ones, then 50 zeros *) ]} *) @@ -57,8 +55,7 @@ val pad_center : module Convert : sig (** Frequency and amplitude scale conversion functions. *) - val mel_to_hz : - ?htk:bool -> (float, 'a, 'dev) Rune.t -> (float, 'a, 'dev) Rune.t + val mel_to_hz : ?htk:bool -> (float, 'a) Nx.t -> (float, 'a) Nx.t (** [mel_to_hz mel_values] converts mel-scale values to Hz. @param htk Use HTK mel scale formula (default: false, uses Slaney formula) @@ -68,13 +65,12 @@ module Convert : sig {3 Example} {[ - let mel_freqs = Rune.linspace Rune.float32 0.0 2595.0 100 in + let mel_freqs = Nx.linspace Nx.float32 0.0 2595.0 100 in let hz_freqs = Utils.Convert.mel_to_hz mel_freqs in (* Convert 100 mel frequencies to Hz *) ]} *) - val hz_to_mel : - ?htk:bool -> (float, 'a, 'dev) Rune.t -> (float, 'a, 'dev) Rune.t + val hz_to_mel : ?htk:bool -> (float, 'a) Nx.t -> (float, 'a) Nx.t (** [hz_to_mel hz_values] converts Hz frequency values to mel scale. @param htk Use HTK mel scale formula (default: false, uses Slaney formula) @@ -84,17 +80,17 @@ module Convert : sig This is the inverse function of {!mel_to_hz}. *) (** Reference value specification for dB conversions. *) - type ('a, 'dev) reference = + type 'a reference = | RefFloat of float (** Use a constant reference value. *) - | RefFunction of ((float, 'a, 'dev) Rune.t -> float) + | RefFunction of ((float, 'a) Nx.t -> float) (** Compute reference from the input tensor (e.g., maximum value). *) val power_to_db : ?amin:float -> ?top_db:float - -> ('a, 'dev) reference - -> (float, 'a, 'dev) Rune.t - -> (float, 'a, 'dev) Rune.t + -> 'a reference + -> (float, 'a) Nx.t + -> (float, 'a) Nx.t (** [power_to_db reference power_spectrum] converts power values to decibels. @param amin Minimum amplitude threshold (default: 1e-10) @@ -110,16 +106,13 @@ module Convert : sig {[ let power_spec = (* compute power spectrum *) in let db_spec = Utils.Convert.power_to_db - (RefFunction Rune.max) + (RefFunction (fun x -> Nx.item [] (Nx.max x))) power_spec in (* Convert to dB relative to maximum value *) ]} *) val db_to_power : - ?amin:float - -> ('a, 'dev) reference - -> (float, 'a, 'dev) Rune.t - -> (float, 'a, 'dev) Rune.t + ?amin:float -> 'a reference -> (float, 'a) Nx.t -> (float, 'a) Nx.t (** [decibels_to_power reference db_values] converts decibel values to power. This is the inverse function of {!power_to_db}. @@ -137,10 +130,9 @@ val melfreqs : -> ?f_min:float -> ?f_max:float -> ?htk:bool - -> 'dev Rune.device - -> (float, 'b) Rune.dtype - -> (float, 'b, 'dev) Rune.t -(** [melfreqs device dtype] generates mel-scale frequency values. + -> (float, 'b) Nx.dtype + -> (float, 'b) Nx.t +(** [melfreqs dtype] generates mel-scale frequency values. Computes a tensor of acoustic frequencies tuned to the mel scale, which approximates human auditory perception. This is used internally by @@ -150,7 +142,6 @@ val melfreqs : @param f_min Minimum frequency in Hz (default: 0.0) @param f_max Maximum frequency in Hz (default: 11025.0) @param htk Use HTK mel scale formula (default: false) - @param device Rune device on which the Tensor should be created @param dtype Data type for the output tensor @return Tensor of mel-scale frequencies @@ -161,10 +152,10 @@ val melfreqs : {[ let mel_freqs = Utils.melfreqs - Rune.float32 ~n_mels:80 ~f_min:80.0 - ~f_max:7600.0 in + ~f_max:7600.0 + Nx.float32 in (* 80 mel-scale frequencies from 80 Hz to 7600 Hz *) ]} @@ -173,17 +164,17 @@ val melfreqs : (** {2 Mathematical Utilities} *) val outer : - (('a, 'b, 'dev) Rune.t -> ('a, 'b, 'dev) Rune.t -> ('a, 'b, 'dev) Rune.t) - -> ('a, 'b, 'dev) Rune.t - -> ('a, 'b, 'dev) Rune.t - -> ('a, 'b, 'dev) Rune.t + (('a, 'b) Nx.t -> ('a, 'b) Nx.t -> ('a, 'b) Nx.t) + -> ('a, 'b) Nx.t + -> ('a, 'b) Nx.t + -> ('a, 'b) Nx.t (** [outer_product operation x y] computes generalized outer product. Applies a binary operation between all pairs of elements from two tensors, creating a matrix where result[i,j] = operation(x[i], y[j]). This is useful for creating distance matrices, correlation matrices, and other pairwise operations. - @param operation Binary operation to apply (e.g., Rune.add, Rune.mul) + @param operation Binary operation to apply (e.g., Nx.add, Nx.mul) @param x First input tensor (1D) @param y Second input tensor (1D) @return 2D tensor with shape [length(x); length(y)] @@ -193,8 +184,8 @@ val outer : {3 Example} {[ - let x = Rune.arange Rune.float32 1.0 4.0 1.0 in - let y = Rune.arange Rune.float32 1.0 3.0 1.0 in - let product = Utils.outer_product Rune.mul x y in + let x = Nx.arange_f Nx.float32 1.0 4.0 1.0 in + let y = Nx.arange_f Nx.float32 1.0 3.0 1.0 in + let product = Utils.outer Nx.mul x y in (* Multiplication table: [[1,2], [2,4], [3,6]] *) ]} *) diff --git a/src/window.ml b/src/window.ml index 253f990..ccd5010 100644 --- a/src/window.ml +++ b/src/window.ml @@ -21,67 +21,60 @@ type window = [`Hanning | `Hamming | `Blackman | `Boxcar] -let cosine_sum ?(fftbins = false) (device : 'dev Rune.device) - (dtype : (float, 'b) Rune.dtype) (a : float array) m = +let cosine_sum ?(fftbins = false) (dtype : (float, 'b) Nx.dtype) + (a : float array) m = if m < 0 then invalid_arg "Window length M must be a non-negative integer" - else if m = 0 then Rune.empty device dtype [|0|] - else if m = 1 then Rune.ones device dtype [|1|] + else if m = 0 then Nx.empty dtype [|0|] + else if m = 1 then Nx.ones dtype [|1|] else let sym = not fftbins in let fac = - if sym then Rune.linspace device dtype (-.Float.pi) Float.pi m + if sym then Nx.linspace dtype (-.Float.pi) Float.pi m else - let full_range = - Rune.linspace device dtype (-.Float.pi) Float.pi (m + 1) - in - Rune.slice [Rs (0, m, 1)] full_range + let full_range = Nx.linspace dtype (-.Float.pi) Float.pi (m + 1) in + Nx.slice [Rs (0, m, 1)] full_range in let w = - if Array.length a > 0 then Rune.full device dtype [|m|] a.(0) - else Rune.zeros device dtype [|m|] + if Array.length a > 0 then Nx.full dtype [|m|] a.(0) + else Nx.zeros dtype [|m|] in let w = Array.fold_left (fun acc (k, coeff) -> if k = 0 || coeff = 0.0 then acc else - let cos_args = Rune.mul_s fac (float_of_int k) in - let cos_terms = Rune.cos cos_args in - let term = Rune.mul_s cos_terms coeff in - Rune.add acc term ) + let cos_args = Nx.mul_s fac (float_of_int k) in + let cos_terms = Nx.cos cos_args in + let term = Nx.mul_s cos_terms coeff in + Nx.add acc term ) w (Array.mapi (fun i x -> (i, x)) a) in w -let hanning ?(fftbins = false) (device : 'dev Rune.device) - (dtype : (float, 'b) Rune.dtype) m = - cosine_sum ~fftbins device dtype [|0.5; 1. -. 0.5|] m +let hanning ?(fftbins = false) (dtype : (float, 'b) Nx.dtype) m = + cosine_sum ~fftbins dtype [|0.5; 1. -. 0.5|] m -let hamming ?(fftbins = false) (device : 'dev Rune.device) - (dtype : (float, 'b) Rune.dtype) m = - cosine_sum ~fftbins device dtype [|0.54; 1. -. 0.54|] m +let hamming ?(fftbins = false) (dtype : (float, 'b) Nx.dtype) m = + cosine_sum ~fftbins dtype [|0.54; 1. -. 0.54|] m -let blackman ?(fftbins = false) (device : 'dev Rune.device) - (dtype : (float, 'b) Rune.dtype) m = - cosine_sum ~fftbins device dtype [|0.42; 0.5; 0.08|] m +let blackman ?(fftbins = false) (dtype : (float, 'b) Nx.dtype) m = + cosine_sum ~fftbins dtype [|0.42; 0.5; 0.08|] m -let boxcar (device : 'dev Rune.device) (dtype : ('a, 'b) Rune.dtype) (size : int) - : ('a, 'b, 'dev) Rune.t = +let boxcar (dtype : ('a, 'b) Nx.dtype) (size : int) : ('a, 'b) Nx.t = if size < 0 then failwith "Window length M must be non-negative" - else if size = 0 then Rune.empty device dtype [|0|] - else Rune.ones device dtype [|size|] + else if size = 0 then Nx.empty dtype [|0|] + else Nx.ones dtype [|size|] -let get (typ : window) (device : 'dev Rune.device) - (dtype : (float, 'b) Rune.dtype) : - ?fftbins:bool -> int -> (float, 'b, 'dev) Rune.t = +let get (typ : window) (dtype : (float, 'b) Nx.dtype) : + ?fftbins:bool -> int -> (float, 'b) Nx.t = fun ?fftbins size -> match typ with | `Hanning -> - hanning ?fftbins device dtype size + hanning ?fftbins dtype size | `Hamming -> - hamming ?fftbins device dtype size + hamming ?fftbins dtype size | `Blackman -> - blackman ?fftbins device dtype size + blackman ?fftbins dtype size | `Boxcar -> - boxcar device dtype size + boxcar dtype size diff --git a/src/window.mli b/src/window.mli index 6691666..f48ca6b 100644 --- a/src/window.mli +++ b/src/window.mli @@ -43,20 +43,14 @@ type window = ] val get : - window - -> 'dev Rune.device - -> (float, 'b) Rune.dtype - -> ?fftbins:bool - -> int - -> (float, 'b, 'dev) Rune.t + window -> (float, 'b) Nx.dtype -> ?fftbins:bool -> int -> (float, 'b) Nx.t (** [get window dtype n] generates a window of size [n] using the specified window function. Creates a window function suitable for spectral analysis applications. The window is normalized and ready to use with STFT or other windowed operations. @param window Window function type to generate - @param device Device on which the resulting Tensor should be created - @param dtype Data type for the window values (e.g., Rune.float32, Rune.float64) + @param dtype Data type for the window values (e.g., Nx.float32, Nx.float64) @param fftbins Whether to use FFT-compatible window length (default: true) @param n Window size in samples (must be > 0) @return Window function tensor of length n @@ -67,13 +61,13 @@ val get : Basic window generation: {[ - let hann_window = Window.get `Hanning Rune.float32 1024 in + let hann_window = Window.get `Hanning Nx.float32 1024 in (* 1024-sample Hann window *) ]} Window for STFT: {[ - let window = Window.get `Hamming Rune.float64 512 in + let window = Window.get `Hamming Nx.float64 512 in (* Use with Transform.stft ~window:(Custom_window window) *) ]} @@ -86,11 +80,10 @@ val get : val cosine_sum : ?fftbins:bool - -> 'dev Rune.device - -> (float, 'b) Rune.dtype + -> (float, 'b) Nx.dtype -> float array -> int - -> (float, 'b, 'dev) Rune.t + -> (float, 'b) Nx.t (** [cosine_sum dtype coeffs n] generates a generalized cosine-sum window. Creates a window function as a weighted sum of cosine terms. This is the @@ -98,8 +91,7 @@ val cosine_sum : windows as special cases. Useful for creating custom window functions. @param fftbins Whether to use FFT-compatible window (default: true) - @param device Device on which the resulting Tensor should be created - @param dtype Data type for window values (e.g., Rune.float32, Rune.float64) + @param dtype Data type for window values (e.g., Nx.float32, Nx.float64) @param coeffs Coefficients for the cosine terms (must have length >= 1) @param n Window size in samples (must be > 0) @return Cosine-sum window tensor of length n @@ -112,14 +104,14 @@ val cosine_sum : Create a Hann window using cosine_sum: {[ let hann_coeffs = [|0.5; -0.5|] in - let hann_window = Window.cosine_sum Rune.float32 hann_coeffs 1024 in - (* Equivalent to Window.hanning Rune.float32 1024 *) + let hann_window = Window.cosine_sum Nx.float32 hann_coeffs 1024 in + (* Equivalent to Window.hanning Nx.float32 1024 *) ]} Create a custom window: {[ let custom_coeffs = [|0.4; -0.4; 0.2|] in - let custom_window = Window.cosine_sum Rune.float32 custom_coeffs 512 in + let custom_window = Window.cosine_sum Nx.float32 custom_coeffs 512 in (* Custom 3-term cosine window *) ]} @@ -128,12 +120,7 @@ val cosine_sum : See: {{:https://en.wikipedia.org/wiki/Window_function#Cosine-sum_windows}Wikipedia: Cosine-sum windows} See: {{:https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.windows.general_cosine.html}SciPy general_cosine} *) -val hanning : - ?fftbins:bool - -> 'dev Rune.device - -> (float, 'b) Rune.dtype - -> int - -> (float, 'b, 'dev) Rune.t +val hanning : ?fftbins:bool -> (float, 'b) Nx.dtype -> int -> (float, 'b) Nx.t (** [hanning dtype n] generates a Hann window of size [n]. The Hann window (also called Hanning window) is a raised cosine window @@ -141,8 +128,7 @@ val hanning : It offers a good balance between frequency resolution and spectral leakage. @param fftbins Whether to use FFT-compatible window (default: true) - @param device Device on which the resulting Tensor should be created - @param dtype Data type for window values (e.g., Rune.float32, Rune.float64) + @param dtype Data type for window values (e.g., Nx.float32, Nx.float64) @param n Window size in samples (must be > 0) @return Hann window tensor of length n @@ -151,18 +137,13 @@ val hanning : {3 Example} {[ - let window = Window.hanning Rune.float32 1024 in + let window = Window.hanning Nx.float32 1024 in (* Use with STFT: Transform.stft ~window:(Custom_window window) signal *) ]} Mathematical definition: w(n) = 0.5 * (1 - cos(2π*n/(N-1))) *) -val hamming : - ?fftbins:bool - -> 'dev Rune.device - -> (float, 'b) Rune.dtype - -> int - -> (float, 'b, 'dev) Rune.t +val hamming : ?fftbins:bool -> (float, 'b) Nx.dtype -> int -> (float, 'b) Nx.t (** [hamming dtype n] generates a Hamming window of size [n]. The Hamming window is similar to the Hann window but uses different @@ -170,8 +151,7 @@ val hamming : the cost of a slightly wider main lobe. @param fftbins Whether to use FFT-compatible window (default: true) - @param device Device on which the resulting Tensor should be created - @param dtype Data type for window values (e.g., Rune.float32, Rune.float64) + @param dtype Data type for window values (e.g., Nx.float32, Nx.float64) @param n Window size in samples (must be > 0) @return Hamming window tensor of length n @@ -179,12 +159,7 @@ val hamming : Mathematical definition: w(n) = 0.54 - 0.46 * cos(2π*n/(N-1)) *) -val blackman : - ?fftbins:bool - -> 'dev Rune.device - -> (float, 'b) Rune.dtype - -> int - -> (float, 'b, 'dev) Rune.t +val blackman : ?fftbins:bool -> (float, 'b) Nx.dtype -> int -> (float, 'b) Nx.t (** [blackman dtype n] generates a Blackman window of size [n]. The Blackman window provides excellent sidelobe suppression (very low @@ -192,8 +167,7 @@ val blackman : Use when you need minimal spectral leakage. @param fftbins Whether to use FFT-compatible window (default: true) - @param device Device on which the resulting Tensor should be created - @param dtype Data type for window values (e.g., Rune.float32, Rune.float64) + @param dtype Data type for window values (e.g., Nx.float32, Nx.float64) @param n Window size in samples (must be > 0) @return Blackman window tensor of length n @@ -201,15 +175,13 @@ val blackman : Mathematical definition: w(n) = 0.42 - 0.5*cos(2π*n/(N-1)) + 0.08*cos(4π*n/(N-1)) *) -val boxcar : - 'dev Rune.device -> ('a, 'b) Rune.dtype -> int -> ('a, 'b, 'dev) Rune.t -(** [boxcar device dtype n] generates a rectangular (boxcar) window of size [n]. +val boxcar : ('a, 'b) Nx.dtype -> int -> ('a, 'b) Nx.t +(** [boxcar dtype n] generates a rectangular (boxcar) window of size [n]. The boxcar window is simply a rectangular window (all ones) that provides the best frequency resolution but the worst spectral leakage. Use only when you specifically need maximum frequency resolution. - @param device Device on which the resulting Tensor should be created @param dtype Data type for window values @param n Window size in samples (must be > 0) @return Rectangular window tensor of length n (all ones) diff --git a/test/dune b/test/dune index 9134e32..e2ab2bc 100644 --- a/test/dune +++ b/test/dune @@ -3,6 +3,7 @@ (tests (names test_aformat + test_transform test_write test_utils ;test_time @@ -10,13 +11,20 @@ ;test_filter ;test_filterbank ) - (libraries alcotest soundml rune) + (modules + test_aformat + test_transform + test_write + test_utils + test_window + tutils) + (libraries alcotest nx soundml) (action (run %{test}))) ;(tests ; (names test_stft) -; (libraries alcotest str yojson nx rune nx.io soundml) +; (libraries alcotest str yojson nx nx.io soundml) ; (package soundml) ; (deps ; generate_vectors.py diff --git a/test/test_stft.ml b/test/test_stft.ml index ef60584..4266f3b 100644 --- a/test/test_stft.ml +++ b/test/test_stft.ml @@ -37,14 +37,13 @@ let string_to_window = function module StftTestable = struct type a = Complex.t - type b = Rune.complex64_elt + type b = Nx.complex64_elt - let dtype = Rune.Complex64 + let dtype = Nx.Complex128 let typ = "stft" - let generate (_ : (a, b) Rune.dtype) (case : string * string * Parameters.t) - (audio : (float, Bigarray.float64_elt, [`c]) Rune.t) = + let generate (case : string * string * Parameters.t) (audio : Nx.float64_t) = let _, _, params = case in let n_fft = Option.value ~default:2048 @@ Parameters.get_int "n_fft" params @@ -53,20 +52,17 @@ module StftTestable = struct Option.value ~default:512 @@ Parameters.get_int "hop_size" params in let win_length = - Option.value ~default:2048 @@ Parameters.get_int "window_length" params + Option.value ~default:n_fft @@ Parameters.get_int "window_length" params in let window = string_to_window - (Option.value ~default:"hann" @@ Parameters.get_string "window" params) + ( Option.value ~default:"hann" + @@ Parameters.get_string "window_type" params ) in let center = Option.value ~default:false @@ Parameters.get_bool "center" params in - let stft = - Transform.stft ~n_fft ~hop_length:hop_size ~win_length ~window ~center - audio - in - stft + Transform.stft ~n_fft ~hop_length:hop_size ~win_length ~window ~center audio end module Tests = Tests_cases (StftTestable) diff --git a/test/test_timeseries.ml b/test/test_timeseries.ml index 107747b..649a052 100644 --- a/test/test_timeseries.ml +++ b/test/test_timeseries.ml @@ -24,15 +24,14 @@ open Vutils module Timeseries = struct type a = float - type b = Rune.float64_elt + type b = Nx.float64_elt - let dtype = Rune.Float64 + let dtype = Nx.Float64 let typ = "timeseries" - let generate (_ : (a, b) Rune.dtype) (_ : string * string * Parameters.t) - (audio : (float, Bigarray.float64_elt, [`ocaml]) Rune.t) = - Rune.cast dtype audio + let generate (_ : string * string * Parameters.t) (audio : Nx.float64_t) = + audio end module Tests = Tests_cases (Timeseries) diff --git a/test/test_transform.ml b/test/test_transform.ml new file mode 100644 index 0000000..7ad37ae --- /dev/null +++ b/test/test_transform.ml @@ -0,0 +1,39 @@ +open Soundml + +let test_stft_defaults_window_length_to_n_fft () = + let signal = Nx.arange_f Nx.float64 0. 64. 1. in + let stft = Transform.stft ~n_fft:16 ~hop_length:8 signal in + Alcotest.check + (Alcotest.array Alcotest.int) + "STFT shape" [|9; 9|] (Nx.shape stft) + +let test_stft_accepts_float32 () = + let signal = Nx.create Nx.float32 [|4|] [|0.; 1.; 2.; 3.|] in + let actual = + Transform.stft ~window:`Boxcar ~center:false ~n_fft:4 ~hop_length:4 signal + |> Nx.to_array + in + let expected = + [| {Complex.re= 6.; im= 0.} + ; {Complex.re= -2.; im= 2.} + ; {Complex.re= -2.; im= 0.} |] + in + Alcotest.check Alcotest.int "spectrum length" (Array.length expected) + (Array.length actual) ; + Array.iteri + (fun i expected -> + Alcotest.check (Alcotest.float 1e-5) + (Printf.sprintf "bin %d real" i) + expected.Complex.re actual.(i).Complex.re ; + Alcotest.check (Alcotest.float 1e-5) + (Printf.sprintf "bin %d imaginary" i) + expected.Complex.im actual.(i).Complex.im ) + expected + +let () = + Alcotest.run "SoundML Transform" + [ ( "STFT" + , [ Alcotest.test_case "win_length defaults to n_fft" `Quick + test_stft_defaults_window_length_to_n_fft + ; Alcotest.test_case "accepts float32 input" `Quick + test_stft_accepts_float32 ] ) ] diff --git a/test/test_utils.ml b/test/test_utils.ml index a652a47..92c076e 100644 --- a/test/test_utils.ml +++ b/test/test_utils.ml @@ -21,21 +21,20 @@ open Soundml -type data = (float, Rune.float32_elt, [`c]) Rune.t +type data = Nx.float32_t let data_testable : data Alcotest.testable = ( module struct type t = data - let pp : t Fmt.t = Rune.pp + let pp : t Fmt.t = Nx.pp let equal : t -> t -> bool = Tutils.Check.rallclose end ) module Test_pad_center = struct let create_data (arr : float array) : data = - Rune.create Rune.c Rune.float32 [|Array.length arr|] arr - (* Create 1D Rune tensor *) + Nx.create Nx.float32 [|Array.length arr|] arr let test_no_padding () = let input_data = create_data [|1.; 2.; 3.|] in @@ -85,10 +84,7 @@ module Test_pad_center = struct let input_data = create_data [|1.; 2.; 3.; 4.|] in let target_size = 2 in let pad_value = 0. in - let expected_exn = - Invalid_argument - "An error occured while trying to pad: current_size > target_size" - in + let expected_exn = Invalid_argument "size must be >= signal length" in Alcotest.check_raises "error_target_too_small: raises Invalid_argument when target_size < \ input_size" @@ -122,10 +118,7 @@ module Test_pad_center = struct let input_data = create_data [|1.; 2.|] in let target_size = 0 in let pad_value = 0. in - let expected_exn = - Invalid_argument - "An error occured while trying to pad: current_size > target_size" - in + let expected_exn = Invalid_argument "size must be >= signal length" in Alcotest.check_raises "zero_target_non_empty_input: raises Invalid_argument when target_size < \ input_size" @@ -149,7 +142,7 @@ end module Test_melfreq = struct let test_default () = let expected = - Rune.create Rune.c Rune.float32 [|128|] + Nx.create Nx.float32 [|128|] [| 0. ; 26.199787 ; 52.399574 @@ -279,12 +272,12 @@ module Test_melfreq = struct ; 10731.102 ; 11025. |] in - let actual = Utils.melfreqs Tutils.device Rune.float32 in + let actual = Utils.melfreqs Nx.float32 in Alcotest.check data_testable "melfreq_default" expected actual let test_custom () = let expected = - Rune.create Rune.c Rune.float32 [|10|] + Nx.create Nx.float32 [|10|] [| 1000. ; 1203.3604 ; 1431.0475 @@ -297,8 +290,7 @@ module Test_melfreq = struct ; 4000. |] in let actual = - Utils.melfreqs ~n_mels:10 ~f_min:1000. ~f_max:4000. ~htk:true - Tutils.device Rune.float32 + Utils.melfreqs ~n_mels:10 ~f_min:1000. ~f_max:4000. ~htk:true Nx.float32 in Alcotest.check data_testable "melfreq_custom" expected actual @@ -309,23 +301,23 @@ end module Test_outer = struct let test_add () = - let x = Rune.create Rune.c Rune.float32 [|3|] [|1.; 2.; 3.|] in - let y = Rune.create Rune.c Rune.float32 [|4|] [|4.; 5.; 6.; 7.|] in + let x = Nx.create Nx.float32 [|3|] [|1.; 2.; 3.|] in + let y = Nx.create Nx.float32 [|4|] [|4.; 5.; 6.; 7.|] in let expected = - Rune.create Rune.c Rune.float32 [|3; 4|] + Nx.create Nx.float32 [|3; 4|] [|5.; 6.; 7.; 8.; 6.; 7.; 8.; 9.; 7.; 8.; 9.; 10.|] in - let actual = Utils.outer Rune.add x y in + let actual = Utils.outer Nx.add x y in Alcotest.check data_testable "outer_add" expected actual let test_mul () = - let x = Rune.create Rune.c Rune.float32 [|3|] [|1.; 2.; 3.|] in - let y = Rune.create Rune.c Rune.float32 [|4|] [|4.; 5.; 6.; 7.|] in + let x = Nx.create Nx.float32 [|3|] [|1.; 2.; 3.|] in + let y = Nx.create Nx.float32 [|4|] [|4.; 5.; 6.; 7.|] in let expected = - Rune.create Rune.c Rune.float32 [|3; 4|] + Nx.create Nx.float32 [|3; 4|] [|4.; 5.; 6.; 7.; 8.; 10.; 12.; 14.; 12.; 15.; 18.; 21.|] in - let actual = Utils.outer Rune.mul x y in + let actual = Utils.outer Nx.mul x y in Alcotest.check data_testable "outer_mul" expected actual let suite = @@ -336,11 +328,11 @@ end module Test_convert = struct let test_hz_to_mel_htk () = let freqs = - Rune.create Rune.c Rune.float32 [|10|] + Nx.create Nx.float32 [|10|] [|0.; 1225.; 2450.; 3675.; 4900.; 6125.; 7350.; 8575.; 9800.; 11025.|] in let expected = - Rune.create Rune.c Rune.float32 [|10|] + Nx.create Nx.float32 [|10|] [| 0. ; 1140.0684 ; 1695.0864 @@ -357,7 +349,7 @@ module Test_convert = struct let test_mel_to_hz_htk () = let mels = - Rune.create Rune.c Rune.float32 [|10|] + Nx.create Nx.float32 [|10|] [| 0. ; 444.44446 ; 888.8889 @@ -370,7 +362,7 @@ module Test_convert = struct ; 4000. |] in let expected = - Rune.create Rune.c Rune.float32 [|10|] + Nx.create Nx.float32 [|10|] [| 0. ; 338.40695 ; 840.4128 @@ -387,11 +379,11 @@ module Test_convert = struct let test_power_to_db () = let s = - Rune.create Rune.c Rune.float32 [|2; 4|] + Nx.create Nx.float32 [|2; 4|] [|1.; 0.1; 0.01; 0.001; 1e-11; 1e-12; 1e-13; 1e-14|] in let expected_ref1_topdb80 = - Rune.create Rune.c Rune.float32 [|2; 4|] + Nx.create Nx.float32 [|2; 4|] [|0.; -10.; -20.; -30.; -80.; -80.; -80.; -80.|] in let actual_ref1_topdb80 = @@ -400,34 +392,24 @@ module Test_convert = struct Alcotest.check data_testable "power_to_db ref=1.0 top_db=80.0" expected_ref1_topdb80 actual_ref1_topdb80 ; let expected_refmax_topdb80 = - Rune.create Rune.c Rune.float32 [|2; 4|] + Nx.create Nx.float32 [|2; 4|] [|0.; -10.; -20.; -30.; -80.; -80.; -80.; -80.|] in let actual_refmax_topdb80 = Utils.Convert.power_to_db ~top_db:80.0 - (Utils.Convert.RefFunction (fun x -> Rune.item [] (Rune.max x))) + (Utils.Convert.RefFunction (fun x -> Nx.item [] (Nx.max x))) s in Alcotest.check data_testable "power_to_db ref=max top_db=80.0" - expected_refmax_topdb80 actual_refmax_topdb80 ; - let expected_ref1_topdbNone = - Rune.create Rune.c Rune.float32 [|2; 4|] - [|0.; -10.; -20.; -30.; -110.; -120.; -130.; -140.|] - in - let actual_ref1_topdbNone = - Utils.Convert.power_to_db ~amin:1e-10 ?top_db:None - (Utils.Convert.RefFloat 1.0) s - in - Alcotest.check data_testable "power_to_db ref=1.0 top_db=None" - expected_ref1_topdbNone actual_ref1_topdbNone + expected_refmax_topdb80 actual_refmax_topdb80 let test_db_to_power () = let db_s = - Rune.create Rune.c Rune.float32 [|2; 4|] + Nx.create Nx.float32 [|2; 4|] [|0.; -10.; -20.; -30.; -80.; -80.; -80.; -80.|] in let expected = - Rune.create Rune.c Rune.float32 [|2; 4|] + Nx.create Nx.float32 [|2; 4|] [|1.; 0.1; 0.01; 0.001; 1e-8; 1e-8; 1e-8; 1e-8|] in let actual = Utils.Convert.db_to_power (Utils.Convert.RefFloat 1.0) db_s in diff --git a/test/test_window.ml b/test/test_window.ml index ddd8061..143709a 100644 --- a/test/test_window.ml +++ b/test/test_window.ml @@ -24,47 +24,36 @@ open Soundml -let f32_nx_testable = - Tutils.tensor_testable Rune.float32 ~rtol:10e-6 ~atol:10e-6 +let f32_nx_testable = Tutils.tensor_testable ~rtol:10e-6 ~atol:10e-6 -let f64_nx_testable = - Tutils.tensor_testable Rune.float64 ~rtol:10e-6 ~atol:10e-6 +let f64_nx_testable = Tutils.tensor_testable ~rtol:10e-6 ~atol:10e-6 let test_m_negative_0 () = Alcotest.check_raises "m_negative: M < 0 should raise Invalid_argument" (Invalid_argument "Window length M must be a non-negative integer") (fun () -> - ignore - (Window.cosine_sum ~fftbins:false Tutils.device Rune.float32 [|1.0|] (-1)) ) + ignore (Window.cosine_sum ~fftbins:false Nx.float32 [|1.0|] (-1)) ) let test_m_zero_f32_1 () = - let expected = Rune.zeros Rune.c Rune.float32 [|0|] in - let actual = - Window.cosine_sum ~fftbins:false Tutils.device Rune.float32 [|1.0|] 0 - in + let expected = Nx.zeros Nx.float32 [|0|] in + let actual = Window.cosine_sum ~fftbins:false Nx.float32 [|1.0|] 0 in Alcotest.check f32_nx_testable "m_zero_f32: M = 0, empty array" expected actual let test_m_zero_f64_2 () = - let expected = Rune.zeros Rune.c Rune.float64 [|0|] in - let actual = - Window.cosine_sum ~fftbins:true Tutils.device Rune.float64 [||] 0 - in + let expected = Nx.zeros Nx.float64 [|0|] in + let actual = Window.cosine_sum ~fftbins:true Nx.float64 [||] 0 in Alcotest.check f64_nx_testable "m_zero_f64: M = 0, empty array" expected actual let test_m_one_f32_3 () = - let expected = Rune.ones Rune.c Rune.float32 [|1|] in - let actual = - Window.cosine_sum ~fftbins:false Tutils.device Rune.float32 [|1.0|] 1 - in + let expected = Nx.ones Nx.float32 [|1|] in + let actual = Window.cosine_sum ~fftbins:false Nx.float32 [|1.0|] 1 in Alcotest.check f32_nx_testable "m_one_f32: M = 1, ones array" expected actual let test_m_one_f64_varied_coeffs_4 () = - let expected = Rune.ones Rune.c Rune.float64 [|1|] in - let actual = - Window.cosine_sum ~fftbins:true Tutils.device Rune.float64 [|0.5; 0.5|] 1 - in + let expected = Nx.ones Nx.float64 [|1|] in + let actual = Window.cosine_sum ~fftbins:true Nx.float64 [|0.5; 0.5|] 1 in Alcotest.check f64_nx_testable "m_one_f64_varied_coeffs: M = 1, ones array" expected actual @@ -76,10 +65,8 @@ let test_sym_m5_a1_f32_5 () = ; 1.00000000e+00 ; 1.00000000e+00 |] in - let expected = Rune.create Rune.c Rune.float32 [|5|] expected_values in - let actual = - Window.cosine_sum ~fftbins:false Tutils.device Rune.float32 [|1.0|] 5 - in + let expected = Nx.create Nx.float32 [|5|] expected_values in + let actual = Window.cosine_sum ~fftbins:false Nx.float32 [|1.0|] 5 in Alcotest.check f32_nx_testable "sym_m5_a1_f32: M = 5, coeffs_a = [np.float64(1.0)], fftbins = False" expected actual @@ -97,10 +84,8 @@ let test_sym_m10_a_half_half_f64_6 () = ; 1.16977778e-01 ; 0.00000000e+00 |] in - let expected = Rune.create Rune.c Rune.float64 [|10|] expected_values in - let actual = - Window.cosine_sum ~fftbins:false Tutils.device Rune.float64 [|0.5; 0.5|] 10 - in + let expected = Nx.create Nx.float64 [|10|] expected_values in + let actual = Window.cosine_sum ~fftbins:false Nx.float64 [|0.5; 0.5|] 10 in Alcotest.check f64_nx_testable "sym_m10_a_half_half_f64: M = 10, coeffs_a = [np.float64(0.5), \ np.float64(0.5)], fftbins = False" @@ -117,10 +102,8 @@ let test_sym_m8_a_hamm_like_f32_7 () = ; 2.53194690e-01 ; 7.99999982e-02 |] in - let expected = Rune.create Rune.c Rune.float32 [|8|] expected_values in - let actual = - Window.cosine_sum ~fftbins:false Tutils.device Rune.float32 [|0.54; 0.46|] 8 - in + let expected = Nx.create Nx.float32 [|8|] expected_values in + let actual = Window.cosine_sum ~fftbins:false Nx.float32 [|0.54; 0.46|] 8 in Alcotest.check f32_nx_testable "sym_m8_a_hamm_like_f32: M = 8, coeffs_a = [np.float64(0.54), \ np.float64(0.46)], fftbins = False" @@ -136,10 +119,9 @@ let test_sym_m7_a_multi_f64_8 () = ; 8.75000000e-01 ; 8.50000000e-01 |] in - let expected = Rune.create Rune.c Rune.float64 [|7|] expected_values in + let expected = Nx.create Nx.float64 [|7|] expected_values in let actual = - Window.cosine_sum ~fftbins:false Tutils.device Rune.float64 - [|1.0; 0.2; 0.05|] 7 + Window.cosine_sum ~fftbins:false Nx.float64 [|1.0; 0.2; 0.05|] 7 in Alcotest.check f64_nx_testable "sym_m7_a_multi_f64: M = 7, coeffs_a = [np.float64(1.0), np.float64(0.2), \ @@ -155,10 +137,8 @@ let test_sym_m6_a_empty_f32_9 () = ; 0.00000000e+00 ; 0.00000000e+00 |] in - let expected = Rune.create Rune.c Rune.float32 [|6|] expected_values in - let actual = - Window.cosine_sum ~fftbins:false Tutils.device Rune.float32 [||] 6 - in + let expected = Nx.create Nx.float32 [|6|] expected_values in + let actual = Window.cosine_sum ~fftbins:false Nx.float32 [||] 6 in Alcotest.check f32_nx_testable "sym_m6_a_empty_f32: M = 6, coeffs_a = [], fftbins = False" expected actual @@ -170,10 +150,8 @@ let test_per_m5_a1_f32_10 () = ; 1.00000000e+00 ; 1.00000000e+00 |] in - let expected = Rune.create Rune.c Rune.float32 [|5|] expected_values in - let actual = - Window.cosine_sum ~fftbins:true Tutils.device Rune.float32 [|1.0|] 5 - in + let expected = Nx.create Nx.float32 [|5|] expected_values in + let actual = Window.cosine_sum ~fftbins:true Nx.float32 [|1.0|] 5 in Alcotest.check f32_nx_testable "per_m5_a1_f32: M = 5, coeffs_a = [np.float64(1.0)], fftbins = True" expected actual @@ -191,10 +169,8 @@ let test_per_m10_a_half_half_f64_11 () = ; 3.45491503e-01 ; 9.54915028e-02 |] in - let expected = Rune.create Rune.c Rune.float64 [|10|] expected_values in - let actual = - Window.cosine_sum ~fftbins:true Tutils.device Rune.float64 [|0.5; 0.5|] 10 - in + let expected = Nx.create Nx.float64 [|10|] expected_values in + let actual = Window.cosine_sum ~fftbins:true Nx.float64 [|0.5; 0.5|] 10 in Alcotest.check f64_nx_testable "per_m10_a_half_half_f64: M = 10, coeffs_a = [np.float64(0.5), \ np.float64(0.5)], fftbins = True" @@ -211,10 +187,8 @@ let test_per_m8_a_hamm_like_f32_12 () = ; 5.40000021e-01 ; 2.14730874e-01 |] in - let expected = Rune.create Rune.c Rune.float32 [|8|] expected_values in - let actual = - Window.cosine_sum ~fftbins:true Tutils.device Rune.float32 [|0.54; 0.46|] 8 - in + let expected = Nx.create Nx.float32 [|8|] expected_values in + let actual = Window.cosine_sum ~fftbins:true Nx.float32 [|0.54; 0.46|] 8 in Alcotest.check f32_nx_testable "per_m8_a_hamm_like_f32: M = 8, coeffs_a = [np.float64(0.54), \ np.float64(0.46)], fftbins = True" @@ -230,10 +204,9 @@ let test_per_m7_a_multi_f64_13 () = ; 9.99455743e-01 ; 8.64175993e-01 |] in - let expected = Rune.create Rune.c Rune.float64 [|7|] expected_values in + let expected = Nx.create Nx.float64 [|7|] expected_values in let actual = - Window.cosine_sum ~fftbins:true Tutils.device Rune.float64 - [|1.0; 0.2; 0.05|] 7 + Window.cosine_sum ~fftbins:true Nx.float64 [|1.0; 0.2; 0.05|] 7 in Alcotest.check f64_nx_testable "per_m7_a_multi_f64: M = 7, coeffs_a = [np.float64(1.0), np.float64(0.2), \ @@ -249,10 +222,8 @@ let test_per_m6_a_empty_f64_14 () = ; 0.00000000e+00 ; 0.00000000e+00 |] in - let expected = Rune.create Rune.c Rune.float64 [|6|] expected_values in - let actual = - Window.cosine_sum ~fftbins:true Tutils.device Rune.float64 [||] 6 - in + let expected = Nx.create Nx.float64 [|6|] expected_values in + let actual = Window.cosine_sum ~fftbins:true Nx.float64 [||] 6 in Alcotest.check f64_nx_testable "per_m6_a_empty_f64: M = 6, coeffs_a = [], fftbins = True" expected actual @@ -323,10 +294,8 @@ let test_sym_m64_a_rect_plus_cosine_f32_15 () = ; 7.01490760e-01 ; 6.99999988e-01 |] in - let expected = Rune.create Rune.c Rune.float32 [|64|] expected_values in - let actual = - Window.cosine_sum ~fftbins:false Tutils.device Rune.float32 [|1.0; 0.3|] 64 - in + let expected = Nx.create Nx.float32 [|64|] expected_values in + let actual = Window.cosine_sum ~fftbins:false Nx.float32 [|1.0; 0.3|] 64 in Alcotest.check f32_nx_testable "sym_m64_a_rect_plus_cosine_f32: M = 64, coeffs_a = [np.float64(1.0), \ np.float64(0.3)], fftbins = False" @@ -399,10 +368,8 @@ let test_per_m64_a_rect_plus_cosine_f64_16 () = ; 7.05764416e-01 ; 7.01444582e-01 |] in - let expected = Rune.create Rune.c Rune.float64 [|64|] expected_values in - let actual = - Window.cosine_sum ~fftbins:true Tutils.device Rune.float64 [|1.0; 0.3|] 64 - in + let expected = Nx.create Nx.float64 [|64|] expected_values in + let actual = Window.cosine_sum ~fftbins:true Nx.float64 [|1.0; 0.3|] 64 in Alcotest.check f64_nx_testable "per_m64_a_rect_plus_cosine_f64: M = 64, coeffs_a = [np.float64(1.0), \ np.float64(0.3)], fftbins = True" @@ -416,10 +383,9 @@ let test_sym_m5_a_one_zero_f32_17 () = ; 5.00000000e-01 ; 1.50000000e+00 |] in - let expected = Rune.create Rune.c Rune.float32 [|5|] expected_values in + let expected = Nx.create Nx.float32 [|5|] expected_values in let actual = - Window.cosine_sum ~fftbins:false Tutils.device Rune.float32 - [|1.0; 0.0; 0.5|] 5 + Window.cosine_sum ~fftbins:false Nx.float32 [|1.0; 0.0; 0.5|] 5 in Alcotest.check f32_nx_testable "sym_m5_a_one_zero_f32: M = 5, coeffs_a = [np.float64(1.0), \ diff --git a/test/test_write.ml b/test/test_write.ml index b49fe80..77d961b 100644 --- a/test/test_write.ml +++ b/test/test_write.ml @@ -54,26 +54,30 @@ let file_exists name = let create_test_audio channels samples sample_rate = let shape = if channels > 1 then [|channels; samples|] else [|samples|] in - let data = Rune.zeros Rune.c Rune.float32 shape in - let freq = 23000. in + let data = Nx.zeros Nx.float32 shape in for channel = 0 to channels - 1 do + let frequency = 440. +. (110. *. Float.of_int channel) in for i = 0 to samples - 1 do let idx = if channels > 1 then [channel; i] else [i] in - let value = sin (2. *. Float.pi *. (freq *. Float.of_int channel)) in - let scalar_value = Rune.scalar Rune.c Rune.float32 value in - Rune.set idx scalar_value data + let value = + sin + ( 2. *. Float.pi *. frequency *. Float.of_int i + /. Float.of_int sample_rate ) + in + Nx.set_item idx value data done done ; (data, sample_rate) let create_empty_audio channels sample_rate = let shape = if channels > 1 then [|channels; 0|] else [|0|] in - let data = Rune.zeros Rune.c Rune.float32 shape in + let data = Nx.zeros Nx.float32 shape in (data, sample_rate) let check_write_read name ?(format : Aformat.t = Aformat.{ftype= WAV; sub= PCM_16; endian= FILE}) - channels samples target_sr ext = + ?(rtol = 10e-5) ?(atol = 10e-5) ?max_abs_error channels samples target_sr + ext = let test_name = Printf.sprintf "%s_%dch_%dsamples_%dHz%s" name channels samples target_sr ext @@ -86,20 +90,28 @@ let check_write_read name (file_exists filename) true ; let read_audio, sample_rate = try - Io.read ~mono:(channels = 1) ~sample_rate:target_sr Rune.c - Rune.float32 filename + Io.read ~mono:(channels = 1) ~sample_rate:target_sr Nx.float32 + filename with ex -> Alcotest.failf "Failed to read back file %s: %s" filename (Printexc.to_string ex) in Alcotest.check (Alcotest.array Alcotest.int) - "Channels match after write" (Rune.shape audio) (Rune.shape read_audio) ; + "Channels match after write" (Nx.shape audio) (Nx.shape read_audio) ; Alcotest.check Alcotest.int "Sample rate match after write" target_sr sample_rate ; - Alcotest.check - (Tutils.tensor_testable (Rune.dtype audio) ~rtol:10e-5 ~atol:10e-5) - "Data unchanged after write" audio read_audio ) + match max_abs_error with + | None -> + Alcotest.check + (Tutils.tensor_testable ~rtol ~atol) + "Data unchanged after write" audio read_audio + | Some tolerance -> + let max_error = + Nx.item [] (Nx.max (Nx.abs (Nx.sub audio read_audio))) + in + Alcotest.check (Alcotest.float tolerance) "Maximum absolute error" 0. + max_error ) let check_write_empty name ?(format : Aformat.t = Aformat.{ftype= WAV; sub= PCM_16; endian= FILE}) @@ -110,11 +122,17 @@ let check_write_empty name Alcotest.test_case test_name `Quick (fun () -> let filename = temp_file ~ext test_name in let audio, sr = create_empty_audio channels target_sr in - Alcotest.check - (Alcotest.neg Alcotest.reject) - "Write empty audio don't raise" - (fun () -> ()) - (fun () -> Io.write ~format filename audio sr) ) + Io.write ~format filename audio sr ) + +let read_without_resampling_preserves_sample_rate () = + let filename = temp_file "read_without_resampling_preserves_sample_rate" in + let audio, sample_rate = create_test_audio 1 1024 44100 in + Io.write filename audio sample_rate ; + let _, actual_sample_rate = + Io.read ~res_typ:NONE ~sample_rate:22050 Nx.float32 filename + in + Alcotest.check Alcotest.int "Original sample rate" sample_rate + actual_sample_rate let tests = let wav = Result.get_ok (Aformat.create Aformat.WAV) in @@ -123,15 +141,18 @@ let tests = [ check_write_read "write_f32_mono_wav_deduced" 1 1024 44100 ".wav" ; check_write_read "write_f32_stereo_wav_deduced" 2 1024 44100 ".wav" ; check_write_read "write_f32_stereo_flac_deduced" 2 512 22050 ".flac" - ; check_write_read "write_f32_mono_ogg_deduced" 1 2048 48000 ".ogg" + ; check_write_read "write_f32_mono_ogg_deduced" ~max_abs_error:0.1 1 2048 + 48000 ".ogg" ; check_write_read "write_f32_stereo_wav_explicit" ~format:wav 2 1024 44100 ".wav" ; check_write_read "write_f32_stereo_flac_explicit" ~format:flac 2 512 22050 ".flac" - ; check_write_read "write_f32_mono_ogg_explicit" ~format:ogg 1 2048 48000 - ".ogg" + ; check_write_read "write_f32_mono_ogg_explicit" ~format:ogg + ~max_abs_error:0.1 1 2048 48000 ".ogg" ; check_write_empty "write_f32_mono_empty" 1 44100 ".wav" - ; check_write_empty "write_f32_stereo_empty" 2 44100 ".wav" ] + ; check_write_empty "write_f32_stereo_empty" 2 44100 ".wav" + ; Alcotest.test_case "read without resampling preserves sample rate" `Quick + read_without_resampling_preserves_sample_rate ] let suite = [("Write/Read Roundtrip", tests)] diff --git a/test/tutils.ml b/test/tutils.ml index 6c6b7e8..d44ab3f 100644 --- a/test/tutils.ml +++ b/test/tutils.ml @@ -20,90 +20,72 @@ (*****************************************************************************) module Check = struct - let shape (x : ('a, 'b, 'dev) Rune.t) (y : ('a, 'b2, 'dev) Rune.t) = - let shape_x = Rune.shape x in - let shape_y = Rune.shape y in + let shape (x : ('a, 'b) Nx.t) (y : ('a, 'c) Nx.t) = + let shape_x = Nx.shape x in + let shape_y = Nx.shape y in if Array.length shape_x <> Array.length shape_y then false else Array.for_all2 (fun x y -> x = y) shape_x shape_y - let rallclose ?(rtol = 1e-5) ?(atol = 1e-10) (x : (float, 'b, 'dev) Rune.t) - (y : (float, 'b2, 'dev) Rune.t) : bool = + let rallclose ?(rtol = 1e-5) ?(atol = 1e-10) (x : (float, 'b) Nx.t) + (y : (float, 'b) Nx.t) = if not (shape x y) then false - else if Rune.size x = 0 && Rune.size y = 0 then true + else if Nx.numel x = 0 then true else - (* For now, assume tensors are on compatible devices or handle device - mismatch gracefully *) - let abs_diff = Rune.abs (Rune.sub x y) in - let tolerance = - Rune.add - (Rune.mul_s (Rune.abs y) rtol) - (Rune.scalar (Rune.device y) (Rune.dtype y) atol) - in - let comparison_mask = Rune.less_equal abs_diff tolerance in - let all_close = Rune.all comparison_mask in - Rune.item [] all_close = 1 + let abs_diff = Nx.abs (Nx.sub x y) in + let tolerance = Nx.add_s (Nx.mul_s (Nx.abs y) rtol) atol in + Nx.item [] (Nx.all (Nx.less_equal abs_diff tolerance)) - let callclose : type a dev. + let callclose : type a. ?rtol:float -> ?atol:float - -> (Complex.t, a, dev) Rune.t - -> (Complex.t, a, dev) Rune.t + -> (Complex.t, a) Nx.t + -> (Complex.t, a) Nx.t -> bool = fun ?(rtol = 1e-05) ?(atol = 1e-05) x y -> if not (shape x y) then false - else if Rune.size x = 0 && Rune.size y = 0 then true + else if Nx.numel x = 0 then true else - (* For complex tensors, convert to float64 for comparison *) - let x_float = Rune.cast Rune.float64 x in - let y_float = Rune.cast Rune.float64 y in - let abs_diff = Rune.abs (Rune.sub x_float y_float) in - let abs_y = Rune.abs y_float in - let tolerance = - Rune.add (Rune.mul_s abs_y rtol) - (Rune.scalar (Rune.device abs_y) (Rune.dtype abs_y) atol) - in - let comparison_mask = Rune.less_equal abs_diff tolerance in - let all_close = Rune.all comparison_mask in - Rune.item [] all_close = 1 + Array.for_all2 + (fun actual expected -> + Complex.norm (Complex.sub actual expected) + <= atol +. (rtol *. Complex.norm expected) ) + (Nx.to_array x) (Nx.to_array y) end -let allclose_aux : type a b dev. - (a, b) Rune.dtype +let allclose_aux : type a b. + (a, b) Nx.dtype -> ?rtol:float -> ?atol:float - -> (a, b, dev) Rune.t - -> (a, b, dev) Rune.t + -> (a, b) Nx.t + -> (a, b) Nx.t -> bool = fun kd ?(rtol = 1e-5) ?(atol = 1e-8) x y -> match kd with - | Rune.Complex32 -> + | Nx.Complex64 -> Check.callclose ~rtol ~atol x y - | Rune.Complex64 -> + | Nx.Complex128 -> Check.callclose ~rtol ~atol x y - | Rune.Float16 -> + | Nx.Float16 -> Check.rallclose ~rtol ~atol x y - | Rune.Float32 -> + | Nx.Float32 -> Check.rallclose ~rtol ~atol x y - | Rune.Float64 -> + | Nx.Float64 -> + Check.rallclose ~rtol ~atol x y + | Nx.BFloat16 -> + Check.rallclose ~rtol ~atol x y + | Nx.Float8_e4m3 -> + Check.rallclose ~rtol ~atol x y + | Nx.Float8_e5m2 -> Check.rallclose ~rtol ~atol x y | _ -> failwith "Unsupported datatype." -let allclose : type a b dev. - ?rtol:float - -> ?atol:float - -> (a, b, dev) Rune.t - -> (a, b, dev) Rune.t - -> bool = - fun ?rtol ?atol x y -> allclose_aux (Rune.dtype x) ?rtol ?atol x y +let allclose : type a b. + ?rtol:float -> ?atol:float -> (a, b) Nx.t -> (a, b) Nx.t -> bool = + fun ?rtol ?atol x y -> allclose_aux (Nx.dtype x) ?rtol ?atol x y -let tensor_testable : type a b dev. - (a, b) Rune.dtype - -> rtol:float - -> atol:float - -> (a, b, dev) Rune.t Alcotest.testable = - fun _ ~rtol ~atol -> +let tensor_testable : type a b. + rtol:float -> atol:float -> (a, b) Nx.t Alcotest.testable = + fun ~rtol ~atol -> let equal a b = allclose ~rtol ~atol a b in - Alcotest.testable Rune.pp equal - -let device = Rune.c + Alcotest.testable Nx.pp equal diff --git a/test/tutils.mli b/test/tutils.mli index bb2de24..4b51530 100644 --- a/test/tutils.mli +++ b/test/tutils.mli @@ -23,43 +23,27 @@ *) module Check : sig val rallclose : - ?rtol:float - -> ?atol:float - -> (float, 'b, 'dev) Rune.t - -> (float, 'b, 'dev) Rune.t - -> bool + ?rtol:float -> ?atol:float -> (float, 'b) Nx.t -> (float, 'b) Nx.t -> bool (** Real-valued all-close function *) val callclose : - 'a 'dev. + 'a. ?rtol:float -> ?atol:float - -> (Complex.t, 'a, 'dev) Rune.t - -> (Complex.t, 'a, 'dev) Rune.t + -> (Complex.t, 'a) Nx.t + -> (Complex.t, 'a) Nx.t -> bool (** Complex-valued all-close function *) - val shape : ('a, 'b, 'dev) Rune.t -> ('a, 'b, 'dev) Rune.t -> bool + val shape : ('a, 'b) Nx.t -> ('a, 'c) Nx.t -> bool (** Check the shape of two ndarrays are equal *) end val allclose : - 'a 'b 'dev. - ?rtol:float - -> ?atol:float - -> ('a, 'b, 'dev) Rune.t - -> ('a, 'b, 'dev) Rune.t - -> bool -(** Checks if two Rune tensors are allclose. This is equivalent to + 'a 'b. ?rtol:float -> ?atol:float -> ('a, 'b) Nx.t -> ('a, 'b) Nx.t -> bool +(** Checks if two Nx tensors are allclose. This is equivalent to NumPy's allclose function. *) val tensor_testable : - 'a 'b 'dev. - ('a, 'b) Rune.dtype - -> rtol:float - -> atol:float - -> ('a, 'b, 'dev) Rune.t Alcotest.testable -(** An Alcotest.testable for Rune tensors. *) - -val device : [`c] Rune.device -(** Device used to run the tests *) + 'a 'b. rtol:float -> atol:float -> ('a, 'b) Nx.t Alcotest.testable +(** An Alcotest testable for Nx tensors. *) diff --git a/test/vutils.ml b/test/vutils.ml index f4977e4..652e0ca 100644 --- a/test/vutils.ml +++ b/test/vutils.ml @@ -152,51 +152,22 @@ module type Testable = sig type b - val dtype : (a, b) Rune.dtype + val dtype : (a, b) Nx.dtype val typ : string - val generate : - (a, b) Rune.dtype - -> string * string * Parameters.t - -> (float, Bigarray.float64_elt, [`c]) Rune.t - -> (a, b, [`c]) Rune.t + val generate : string * string * Parameters.t -> Nx.float64_t -> (a, b) Nx.t end module Tests_cases (T : Testable) = struct include T - let read_audio (type c) (audio_dtype : (float, c) Rune.dtype) (path : string) + let read_audio (type c) (audio_dtype : (float, c) Nx.dtype) (path : string) (res_typ : Io.resampling_t) (sample_rate : int) (mono : bool) = - let audio, _ = - Io.read ~res_typ ~sample_rate ~mono Rune.c audio_dtype path - in + let audio, _ = Io.read ~res_typ ~sample_rate ~mono audio_dtype path in audio - let read_npy : type a b. (a, b) Rune.dtype -> string -> (a, b, [`c]) Rune.t = - fun dtype path -> - let packed = Nx_io.load_npy path in - let device = Rune.c in - match dtype with - | Rune.Float32 -> - let nx_tensor = Nx_io.as_float32 packed in - Rune.of_bigarray device (Nx.to_bigarray nx_tensor) - | Rune.Float64 -> - let nx_tensor = Nx_io.as_float64 packed in - Rune.of_bigarray device (Nx.to_bigarray nx_tensor) - | Rune.Complex64 -> - let nx_tensor = Nx_io.as_complex64 packed in - Rune.of_bigarray device (Nx.to_bigarray nx_tensor) - | _ -> - failwith "Unsupported datatype" - (* let packed = Nx_io.load_npy path in let device = Rune.c in match dtype with - | Rune.Float32 -> let nx_tensor = Nx.as_float32 packed in Rune.of_bigarray - device (Nx.to_bigarray nx_tensor) | Rune.Float64 -> let nx_tensor = - Nx_io.to_float64 packed in Rune.of_bigarray device (Nx.to_bigarray - nx_tensor) | Rune.Complex32 -> let nx_tensor = Nx_io.to_complex32 packed in - Rune.of_bigarray device (Nx.to_bigarray nx_tensor) | Rune.Complex64 -> let - nx_tensor = Nx_io.to_complex64 packed in Rune.of_bigarray device - (Nx.to_bigarray nx_tensor) | _ -> failwith "Unsupported datatype" *) + let read_npy dtype path = Nx_io.to_typed dtype (Nx_io.load_npy path) let create_tests (data : (string * string * Parameters.t) list) (rtol : float) (atol : float) : unit Alcotest.test_case list = @@ -221,12 +192,12 @@ module Tests_cases (T : Testable) = struct ( Option.value ~default:"None" @@ Parameters.get_string "res_type" params ) in - let audio = read_audio Rune.float64 audio_path resampler sr mono in - let generated = generate dtype case audio in + let audio = read_audio Nx.float64 audio_path resampler sr mono in + let generated = generate case audio in let vector = read_npy dtype vector_path in let test_dense () = Alcotest.check - (Tutils.tensor_testable dtype ~rtol ~atol) + (Tutils.tensor_testable ~rtol ~atol) (typ ^ "_dense_" ^ basename) generated vector in diff --git a/test/vutils.mli b/test/vutils.mli index aa66fbc..e9175b8 100644 --- a/test/vutils.mli +++ b/test/vutils.mli @@ -68,29 +68,25 @@ module type Testable = sig type b - val dtype : (a, b) Rune.dtype + val dtype : (a, b) Nx.dtype (** Type of the data generated by the functionnality we're testing *) val typ : string - val generate : - (a, b) Rune.dtype - -> string * string * Parameters.t - -> (float, Bigarray.float64_elt, [`c]) Rune.t (* Audio is always float *) - -> (a, b, [`c]) Rune.t + val generate : string * string * Parameters.t -> Nx.float64_t -> (a, b) Nx.t (** Generate a testing vector given some parameters and an audio array *) end module Tests_cases (_ : Testable) : sig val read_audio : - (float, 'a) Rune.dtype + (float, 'a) Nx.dtype -> string -> Io.resampling_t -> int -> bool - -> (float, 'a, [`c]) Rune.t + -> (float, 'a) Nx.t - val read_npy : ('a, 'b) Rune.dtype -> string -> ('a, 'b, [`c]) Rune.t + val read_npy : ('a, 'b) Nx.dtype -> string -> ('a, 'b) Nx.t val create_tests : (string * string * Parameters.t) list