General reading and writing of WAV files

hifiscan/__init__.py

@@ -3,4 +3,5 @@
 from hifiscan.analyzer import (
     Analyzer, XY, geom_chirp, linear_chirp, minimum_phase, resample,
     smooth, taper, tone, window)
-from hifiscan.audio import Audio, read_correction, write_wav
+from hifiscan.audio import Audio
+from hifiscan.io_ import Sound, read_correction, read_wav, write_wav

hifiscan/analyzer.py

@@ -1,12 +1,14 @@
 import array
 import types
 from functools import lru_cache
-from typing import List, NamedTuple, Optional, Tuple
+from typing import NamedTuple, Optional, Tuple
 
 import numpy as np
 from numba import njit
 from numpy.fft import fft, ifft, irfft, rfft
 
+from hifiscan.io_ import Correction
+
 
 class XY(NamedTuple):
     """XY coordinate data of arrays of the same length."""
@@ -15,9 +17,6 @@ class XY(NamedTuple):
     y: np.ndarray
 
 
-Correction = List[Tuple[float, float]]
-
-
 class Analyzer:
     """
     Analyze the system response to a chirp stimulus.

hifiscan/app.py

@@ -149,7 +149,8 @@ class App(qt.QMainWindow):
             self, 'Save inverse impulse response',
             str(self.saveDir / name), 'WAV (*.wav)')
         if filename:
-            hifi.write_wav(filename, analyzer.rate, irInv)
+            data = np.vstack([irInv, irInv])
+            hifi.write_wav(filename, data, analyzer.rate)
             self.saveDir = Path(filename).parent
 
     def run(self):

hifiscan/audio.py

@@ -1,6 +1,4 @@
 import array
-import sys
-import wave
 from collections import deque
 from dataclasses import dataclass
 from typing import AsyncIterator, Deque
@@ -9,8 +7,6 @@ import eventkit as ev
 import numpy as np
 import sounddevice as sd
 
-from hifiscan.analyzer import Correction
-
 
 class Audio:
     """
@@ -91,30 +87,3 @@ class PlayItem:
         return chunk
 
 
-def write_wav(path: str, rate: int, sound: np.ndarray):
-    """
-    Write a 1-channel float array with values between -1 and 1
-    as a 32 bit stereo wave file.
-    """
-    scaling = 2**31 - 1
-    mono = np.asarray(sound * scaling, np.int32)
-    if sys.byteorder == 'big':
-        mono = mono.byteswap()
-    stereo = np.vstack([mono, mono]).flatten(order='F')
-    with wave.open(path, 'wb') as wav:
-        wav.setnchannels(2)
-        wav.setsampwidth(4)
-        wav.setframerate(rate)
-        wav.writeframes(stereo.tobytes())
-
-
-def read_correction(path: str) -> Correction:
-    corr = []
-    with open(path, 'r') as f:
-        for line in f.readlines():
-            try:
-                freq, db = line.split()
-                corr.append((float(freq), float(db)))
-            except ValueError:
-                pass
-    return corr

hifiscan/io_.py

@@ -0,0 +1,88 @@
+import wave
+from typing import List, NamedTuple, Tuple
+
+import numpy as np
+
+
+class Sound(NamedTuple):
+    data: np.ndarray
+    rate: int
+    width: int = 4
+
+
+Correction = List[Tuple[float, float]]
+
+
+def write_wav(path: str, data: np.ndarray, rate: int, width: int = 4):
+    """
+    Write n-channel float array with values between -1 and 1 to WAV file
+
+    Params:
+      path: Filename of WAV file.
+      data: Sound sample data array.
+      rate: Sample rate in Hz.
+      width: Sample width in bytes.
+    """
+    if not width in [1, 2, 3, 4]:
+        raise ValueError(f'Invalid sample width: {width}')
+    data = np.asarray(data)
+    ch = 1 if len(data.shape) < 2 else len(data)
+    if width == 4:
+        arr = np.empty_like(data, np.int32)
+        np.rint((2 ** 31 - 1) * data, out=arr, casting='unsafe')
+    elif width == 3:
+        arr = np.empty_like(data, np.int32)
+        np.rint((2 ** 31 - 2 ** 8) * data, out=arr, casting='unsafe')
+        arr = arr.flatten(order='F').view(np.uint8)
+        # Drop every 4th byte.
+        arr = np.vstack([arr[1::4], arr[2::4], arr[3::4]])
+    elif width == 2:
+        arr = np.empty_like(data, np.int16)
+        np.rint((2 ** 15 - 1) * data, out=arr, casting='unsafe')
+    else:
+        arr = np.empty_like(data, np.int8)
+        np.rint((2 ** 7 - 1) * data, out=arr, casting='unsafe')
+    with wave.open(path, 'wb') as wav:
+        wav.setnchannels(ch)
+        wav.setsampwidth(width)
+        wav.setframerate(rate)
+        wav.writeframes(arr.tobytes(order='F'))
+
+
+def read_wav(path: str) -> Sound:
+    """
+    Read WAV file and return float32 arrays between -1 and 1.
+    """
+    with wave.open(path, 'rb') as wav:
+        ch, width, rate, n, _, _ = wav.getparams()
+        frames = wav.readframes(n)
+    if width == 4:
+        buff = np.frombuffer(frames, np.int32)
+        data = buff.astype('f') / np.float32(2 ** 31 - 1)
+    elif width == 3:
+        buff = np.frombuffer(frames, np.uint8)
+        uints = buff[0::3].astype(np.uint32) << 8 \
+            | buff[1::3].astype(np.uint32) << 16 \
+            | buff[2::3].astype(np.uint32) << 24
+        data = uints.view(np.int32).astype('f') / np.float32(
+            2 ** 31 - 2 ** 8)
+    elif width == 2:
+        buff = np.frombuffer(frames, np.int16)
+        data = buff.astype('f') / np.float32(2 ** 15 - 1)
+    else:
+        buff = np.frombuffer(frames, np.int8)
+        data = buff.astype('f') / np.float32(2 ** 7 - 1)
+    data = data.reshape((-1, ch)).T
+    return Sound(data, rate, width)
+
+
+def read_correction(path: str) -> Correction:
+    corr = []
+    with open(path, 'r') as f:
+        for line in f.readlines():
+            try:
+                freq, db = line.split()
+                corr.append((float(freq), float(db)))
+            except ValueError:
+                pass
+    return corr