src.acoustools.Export.Holo

  1'''
  2Export to .holo file -> List of holograms\n
  3The holo format stores each phase and amplitude as a 5 bit integer with each phase and ampplitude being discritised to 32 levels \n
  4The result of this is then compressed using pythons zlib (https://docs.python.org/3/library/zlib.html)
  5Internally stored as blocks of phase and amplitude seperated by 2^6-1 eg phase1 <11111> amplitude1 <11111> phase2 <11111> amplitude2 <11111> ...
  6'''
  7
  8import torch, io, zlib
  9from acoustools.Utilities.Utilities import batch_list
 10from acoustools.Utilities.Setup import device,DTYPE
 11# from acoustools.Constants import wavelength
 12
 13from torch import Tensor
 14
 15def compress(hologram:Tensor, levels:int=32) -> list[Tensor]:
 16    '''
 17    @private
 18    '''
 19    phase_divs = 2*3.1415/levels
 20    phase_levels = torch.angle(hologram)/phase_divs
 21    phase_levels += levels / 2
 22
 23    amp_divs = 1/levels
 24    amp_levels = torch.abs(hologram)/amp_divs
 25
 26    torch.round_(phase_levels).to(torch.int8)
 27    torch.round_(amp_levels).to(torch.int8)
 28    
 29    return phase_levels, amp_levels
 30
 31def decompress(hologram:list, amplitudes:list, levels:int=32) -> list[Tensor]:
 32    '''
 33    @private
 34    '''
 35    phase_divs = 2*3.1415/levels
 36    phases = [(p - levels/2) * phase_divs for p in hologram]
 37
 38    amp_divs = 1/levels
 39    amplitudes = [a * amp_divs for a in amplitudes]
 40    
 41
 42    holo = [a*torch.e ** (1j*p) for a,p in zip(amplitudes,phases)]
 43    holo = torch.tensor(holo, device=device,dtype=DTYPE).unsqueeze_(0).unsqueeze_(2)
 44    return holo
 45
 46
 47def save_holograms(holos:list[Tensor]|Tensor, fname:str):
 48    '''
 49    Save holograms in .holo format. 
 50    :param holos: Holograms to use
 51    :param fname:  filename to use. Will append .holo is no extension provides
 52    '''
 53    if '.' not in fname:
 54        fname += '.holo'
 55    # pickle.dump(holos, open(fname, 'wb'))
 56    
 57    # with open(fname, 'wb') as file:
 58    with io.BytesIO() as file:
 59        for holo in holos:
 60            phase, amp = compress(holo.squeeze())
 61            for p in phase:
 62                p = int(p.item())
 63                file.write((p).to_bytes(5, byteorder='big', signed=False))
 64            file.write((2**6-1).to_bytes(5, byteorder='big', signed=False))
 65            for a in amp:
 66                a = int(a.item())
 67                file.write((a).to_bytes(5, byteorder='big', signed=False))
 68            file.write((2**6-1).to_bytes(5, byteorder='big', signed=False))
 69        
 70        new_data = zlib.compress(file.getbuffer())
 71        f = open(fname, 'wb')
 72        f.write(new_data)
 73        f.close()
 74
 75def load_holograms(path:str) -> list[Tensor]:
 76    '''
 77    Reads .holo file
 78    :param path: Path to read
 79    :returns holgrams:
 80    '''
 81    if '.' not in path:
 82        path += '.holo'
 83    # holos = pickle.load(open(path, 'rb'))
 84
 85    with open(path, 'rb') as file:
 86        phases = []
 87        amps = []
 88        holos = []
 89        reading_amps = 0
 90        data = file.read()
 91        data = zlib.decompress(data)
 92        for bits in batch_list(data,5):
 93            j = int.from_bytes(bits, byteorder='big')
 94            if j < 2**6-1:
 95                if not reading_amps:
 96                    phases.append(j)
 97                else:
 98                    amps.append(j)
 99            else:
100                if reading_amps == 0:
101                    reading_amps = 1
102                else:
103                    reading_amps = 0
104                    holos.append([phases, amps])
105                    phases = []
106                    amps = []
107        xs = []
108        for h in holos:
109            x = decompress(h[0], h[1])
110            xs.append(x)
111
112    return xs