Entangled units

src/EntangledUnits/EntangledReshaping.jl

@@ -0,0 +1,89 @@
+#TODO: Test this rigorously -- this is key to reshaping EntangledSystems and
+# making EntangledSpinWaveTheorys. 
+
+# An entangled system can be specified with a list of
+# tuples, e.g. [(1,2), (3,4)], which will group the original sites 1 and 2 into
+# one unit and 3 and 4 into another. Given an EntangledSystem constructed from a
+# sys_origin and and a reshaped sys_origin, this function returns a list of
+# tuples for specifying the corresponding entanglement of the reshaped system.
+function units_for_reshaped_system(reshaped_sys_origin, esys)
+    (; sys_origin) = esys                 
+    units = original_unit_spec(esys)
+    new_crystal = reshaped_sys_origin.crystal
+    new_atoms = collect(1:natoms(new_crystal))
+    new_units = []
+
+    # Take a list of all the new atoms in the reshaped system. Pick the first.
+    # Map it back to the original system to determine what unit it belongs to.
+    # Then map all members of the unit forward to define the unit in terms of
+    # the atoms of the reshaped system. Remove these atoms from the list of
+    # sites left to be "entangled" and repeat until list of new atoms is
+    # exhausted.
+    while length(new_atoms) > 0
+        # Pick any site from list of new sites
+        new_atom = new_atoms[1]
+        new_site = CartesianIndex(1, 1, 1, new_atom) # Only need to define for a single unit cell, may as well be the first 
+        new_position = position(reshaped_sys_origin, new_site)
+
+        # Find corresponding original atom number.
+        site = position_to_site(sys_origin, position(reshaped_sys_origin, new_site))
+        original_atom = site[4]
+        position_of_corresponding_atom = position(sys_origin, (1, 1, 1, original_atom))
+        offset = new_position - position_of_corresponding_atom
+
+        # Find the unit to which this original atom belongs.
+        unit = units[findfirst(unit -> original_atom in unit, units)]
+
+        # Find positions of all atoms in the unit, find corresponding sites in reshape system, and define unit for reshaped system.
+        unit_positions = [position(sys_origin, CartesianIndex(1, 1, 1, atom)) for atom in unit]
+        new_unit_sites = [position_to_site(reshaped_sys_origin, position + offset) for position in unit_positions]
+        new_unit = Int64[]
+        for new_site in new_unit_sites
+            i, j, k, a = new_site.I
+            if !(i == j == k == 1)
+                error("Specified reshaping incompatible with specified entangled units. (Unit split between crystalographic unit cells.)")
+            end
+            push!(new_unit, a)
+        end
+        push!(new_units, Tuple(new_unit))
+
+        # Delete members of newly defined unit from list of all atoms in the
+        # reshaped system.
+        idcs = findall(atom -> atom in new_unit, new_atoms)
+        deleteat!(new_atoms, idcs)
+    end
+
+    return new_units
+end
+
+function reshape_supercell(esys::EntangledSystem, shape)
+    (; sys, sys_origin) = esys
+
+    # Reshape the origin System.
+    sys_origin_new = reshape_supercell(sys_origin, shape)
+
+    # Reshape the the underlying "entangled" System.
+    units_new = units_for_reshaped_system(sys_origin_new, esys)
+    _, contraction_info = contract_crystal(sys_origin_new.crystal, units_new)
+    sys_new = reshape_supercell(sys, shape)
+
+    # Construct dipole operator field for reshaped EntangledSystem
+    dipole_operators_origin = all_dipole_observables(sys_origin_new; apply_g=false) 
+    (; observables, source_idcs) = observables_to_product_space(dipole_operators_origin, sys_origin_new, contraction_info)
+
+    return EntangledSystem(sys_new, sys_origin_new, contraction_info, observables, source_idcs)
+end
+
+function repeat_periodically(esys, counts)
+    (; sys, sys_origin, contraction_info) = esys
+
+    # Repeat both entangled and original system periodically
+    sys_new = repeat_periodically(sys, counts)
+    sys_origin_new = repeat_periodically(sys_origin, counts)
+
+    # Construct dipole operator field for reshaped EntangledSystem
+    dipole_operators_origin = all_dipole_observables(sys_origin_new; apply_g=false) 
+    (; observables, source_idcs) = observables_to_product_space(dipole_operators_origin, sys_origin_new, contraction_info)
+
+    return EntangledSystem(sys_new, sys_origin_new, contraction_info, observables, source_idcs)
+end

src/EntangledUnits/EntangledSampledCorrelations.jl

@@ -0,0 +1,170 @@
+struct EntangledSampledCorrelations
+    sc::SampledCorrelations  # Parent SampledCorrelations
+    esys::EntangledSystem    # Probably don't need to carry around the whole thing -- defeats spirit of original design for SC
+end
+
+struct EntangledSampledCorrelationsStatic
+    sc::SampledCorrelationsStatic  # Parent SampledCorrelations
+    esys::EntangledSystem          # Probably don't need to carry around the whole thing -- defeats spirit of original design for SC
+end
+
+function Base.setproperty!(sc::T, sym::Symbol, val) where {T<:Union{EntangledSampledCorrelations, EntangledSampledCorrelationsStatic}}
+    sc = sc.sc
+    if sym == :measure
+        @assert sc.measure.observables ≈ val.observables "New MeasureSpec must contain identical observables."
+        @assert all(x -> x == 1, sc.measure.corr_pairs .== val.corr_pairs) "New MeasureSpec must contain identical correlation pairs."
+        setfield!(sc, :measure, val)
+    else
+        setfield!(sc, sym, val)
+    end
+end
+
+function Base.show(io::IO, ::EntangledSampledCorrelations)
+    print(io, "EntangledSampledCorrelations")
+end
+
+function Base.show(io::IO, ::EntangledSampledCorrelationsStatic)
+    print(io, "EntangledSampledCorrelationsStatic")
+end
+
+function Base.show(io::IO, ::MIME"text/plain", esc::EntangledSampledCorrelations)
+    (; crystal, sys_dims, nsamples) = esc.sc
+    printstyled(io, "EntangledSampledCorrelations"; bold=true, color=:underline)
+    println(io," ($(Base.format_bytes(Base.summarysize(esc))))")
+    print(io,"[")
+    printstyled(io,"S(q,ω)"; bold=true)
+    print(io," | nω = $(round(Int, size(esc.sc.data)[7]/2 + 1)), Δω = $(round(esc.sc.Δω, digits=4))")
+    println(io," | $nsamples $(nsamples > 1 ? "samples" : "sample")]")
+    println(io,"Lattice: $sys_dims × $(natoms(crystal))")
+end
+
+function Base.show(io::IO, ::MIME"text/plain", esc::EntangledSampledCorrelationsStatic)
+    (; crystal, sys_dims, nsamples) = esc.sc.parent
+    printstyled(io, "SampledCorrelationsStatic"; bold=true, color=:underline)
+    println(io," ($(Base.format_bytes(Base.summarysize(esc))))")
+    print(io,"[")
+    printstyled(io,"S(q)"; bold=true)
+    println(io," | $nsamples $(nsamples > 1 ? "samples" : "sample")]")
+    println(io,"Lattice: $sys_dims × $(natoms(crystal))")
+end
+
+function Base.setproperty!(esc::EntangledSampledCorrelations, sym::Symbol, val)
+    if sym == :measure
+        measure = val
+        (; observables) = observables_to_product_space(measure.observables, esc.esys.sys_origin, esc.esys.contraction_info)
+        @assert esc.sc.measure.observables ≈ observables "New MeasureSpec must contain identical observables."
+        @assert all(x -> x == 1, esc.sc.measure.corr_pairs .== measure.corr_pairs) "New MeasureSpec must contain identical correlation pairs."
+        setfield!(esc.sc, :measure, val) # Sets new combiner
+    else
+        setfield!(sc, sym, val)
+    end
+end
+
+Base.setproperty!(esc::EntangledSampledCorrelationsStatic, sym::Symbol, val) = setproperty!(esc.parent, sym, val)
+
+
+# Take observables specified in terms or original system and transform them into
+# a field of observables in the tensor product space, together with mapping
+# information for populating the expectation values of these operators with
+# respect to the entangled system.
+function observables_to_product_space(observables, sys_origin, contraction_info)
+    Ns_per_unit = Ns_in_units(sys_origin, contraction_info)
+    Ns_all = [prod(Ns) for Ns in Ns_per_unit]
+    N = Ns_all[1]
+    @assert all(==(N), Ns_all) "All entangled units must have the same dimension Hilbert space."
+
+    observables_new = fill(zeros(ComplexF64, N, N), size(observables)) 
+    positions = zeros(Vec3, size(observables)[2:end])
+    source_idcs = zeros(Int64, size(observables)[2:end])
+
+    for site in eachsite(sys_origin)
+        atom = site.I[4]
+        positions[site] = sys_origin.crystal.positions[atom]
+        source_idcs[site] = contraction_info.forward[atom][1]
+        for μ in axes(observables, 1)
+            obs = observables[μ, site]
+            unit, unitsub = contraction_info.forward[atom]
+            Ns = Ns_per_unit[unit]
+            observables_new[μ, site] = local_op_to_product_space(obs, unitsub, Ns)
+        end
+    end
+    observables = observables_new
+
+    return (; observables, positions, source_idcs)
+end
+
+
+# Configures an ordinary SampledCorrelations for an entangled system. The
+# measure is assumed to correspond to the sites of the original "unentangled"
+# system. 
+function SampledCorrelations(esys::EntangledSystem; measure, energies, dt, calculate_errors=false)
+    # Convert observables on different sites to "multiposition" observables in
+    # tensor product spaces. With entangled units, the position of an operator
+    # cannot be uniquely determined from the `atom` index of a `Site`. Instead,
+    # the position is recorded independently, and the index of the relevant
+    # coherent state (which may now be used for operators corresponding to
+    # multiple positions) is recorded in `source_idcs`.
+    (; observables, positions, source_idcs) = observables_to_product_space(measure.observables, esys.sys_origin, esys.contraction_info)
+
+    # Make a sampled correlations for the esys.
+    sc = SampledCorrelations(esys.sys; measure, energies, dt, calculate_errors, positions) 
+
+    # Replace relevant fields or the resulting SampledCorrelations. Note use of
+    # undocumented `positions` keyword. This can be eliminated if positions are
+    # migrated into the MeasureSpec.
+    crystal = esys.sys_origin.crystal
+    origin_crystal = orig_crystal(esys.sys_origin)
+    sc_new = SampledCorrelations(sc.data, sc.M, crystal, origin_crystal, sc.Δω, measure, observables, positions, source_idcs, measure.corr_pairs,
+                                 sc.measperiod, sc.dt, sc.nsamples, sc.samplebuf, sc.corrbuf, sc.space_fft!, sc.time_fft!, sc.corr_fft!, sc.corr_ifft!)
+
+    return EntangledSampledCorrelations(sc_new, esys)
+end
+
+function SampledCorrelationsStatic(esys::EntangledSystem; measure, calculate_errors=false)
+    (; observables, positions, source_idcs) = observables_to_product_space(measure.observables, esys.sys_origin, esys.contraction_info)
+    sc = SampledCorrelations(esys.sys; measure, energies=nothing, dt=NaN, calculate_errors, positions) 
+
+    # Replace relevant fields
+    crystal = esys.sys_origin.crystal
+    origin_crystal = orig_crystal(esys.sys_origin)
+    parent = SampledCorrelations(sc.data, sc.M, crystal, origin_crystal, sc.Δω, measure, observables, positions, source_idcs, measure.corr_pairs,
+                                 sc.measperiod, sc.dt, sc.nsamples, sc.samplebuf, sc.corrbuf, sc.space_fft!, sc.time_fft!, sc.corr_fft!, sc.corr_ifft!)
+
+    ssc = SampledCorrelationsStatic(parent)
+    return EntangledSampledCorrelationsStatic(ssc, esys)
+end
+
+
+function add_sample!(esc::EntangledSampledCorrelations, esys::EntangledSystem; window=:cosine)
+    new_sample!(esc.sc, esys.sys)
+    accum_sample!(esc.sc; window)
+end
+
+function add_sample!(esc::EntangledSampledCorrelationsStatic, esys::EntangledSystem; window=:cosine)
+    add_sample!(esc.sc, esys.sys)
+end
+
+available_energies(esc::EntangledSampledCorrelations) = available_energies(esc.sc)
+available_wave_vectors(esc::EntangledSampledCorrelations) = available_wave_vectors(esc.sc)
+
+function clone_correlations(esc::EntangledSampledCorrelations; kwargs...)
+    sc = clone_correlations(esc.sc)
+    EntangledSampledCorrelations(sc, esc.esys)
+end
+
+function merge_correlations(escs::Vector{EntangledSampledCorrelations}; kwargs...)
+    sc_merged = merge_correlations([esc.sc for esc in escs])
+    EntangledSampledCorrelations(sc_merged, escs[1].esys)
+end
+
+function intensities(esc::EntangledSampledCorrelations, qpts; kwargs...)
+    intensities(esc.sc, qpts; kwargs...)
+end
+
+function intensities_static(esc::EntangledSampledCorrelations, qpts; kwargs...)
+    intensities_static(esc.sc, qpts; kwargs...)
+end
+
+function intensities_static(esc::EntangledSampledCorrelationsStatic, qpts; kwargs...)
+    intensities_static(esc.sc, qpts; kwargs...)
+end