From 94a38841acde7b322a351923bf6dbf624e1ca7f2 Mon Sep 17 00:00:00 2001
From: Brandon Flores <bsflores@wisc.edu>
Date: Wed, 29 Nov 2023 14:11:52 -0600
Subject: [PATCH] Fixed issues with `DataSpace(x)` default method

---
 CHANGELOG.md    |  4 ++++
 src/lattices.jl | 28 ++++++++++++----------------
 src/traits.jl   | 20 ++++++++++++++++++++
 3 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4a0229ad..1fecc91b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,10 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
   - The Types section of the documentation has been split up into separate sections for lattice
 basis vectors, atoms and crystal representations, and data grids.
 
+### Fixed
+  - The default definition of `Electrum.DataSpace(x)` is now `DataSpace(typeof(x))`, not
+`DataSpace(typeof(basis(x)))`.
+
 ## [0.1.16]: 2023-11-16
 
 ### Added
diff --git a/src/lattices.jl b/src/lattices.jl
index 68431337..9565b964 100644
--- a/src/lattices.jl
+++ b/src/lattices.jl
@@ -175,26 +175,22 @@ For predictable results, use `convert(T, basis(x))` where `T` is the desired typ
 """
 basis(x) = x.basis::LatticeBasis
 
-#---Data space traits------------------------------------------------------------------------------#
-
-DataSpace(::Type{<:LatticeBasis{S,D}}) where {S,D} = S{D}()
-DataSpace(b::LatticeBasis) = DataSpace(typeof(b))
-
+#---Real/reciprocal space traits-------------------------------------------------------------------#
 """
-    Electrum.DataSpace(x) -> DataSpace
+    Electrum.DataSpace(::Type{<:Electrum.LatticeBasis{S,D}}) = S{D}()
 
-Returns a trait that determines whether a data set associated with a crystal is defined in real
-space (`ByRealSpace{D}()`), reciprocal space (`ByReciprocalSpace{D}()`), or by atomic positions
-(`ByAtom{D}`), where `D` is the number of dimensions.
+Returns the `ByRealSpace{D}()` or `ByReciprocalSpace{D}()` trait objects depending on the type of
+lattice.
 
-By default, `DataSpace(x)` will infer the appropriate trait from the lattice basis vectors
-included in `x`. The fallback definition is:
-
-    DataSpace(x) = DataSpace(typeof(basis(x))) # basis(x) falls back to x.basis
-    DataSpace(::Type{T}) where T = DataSpace(fieldtype(T, :basis))
+Note that this function only works correctly when S === `ByRealSpace` or `ByReciprocalSpace`, and
+not if they are `ByRealSpace{D}` or `ByReciprocalSpace{D}`, or their singleton instances. Even
+though `Electrum.LatticeBasis{ByRealSpace{D},D,T}` is a valid type that has otherwise correct
+behavior, this function only works correctly when for `Electrum.LatticeBasis{ByRealSpace,D}` and
+`Electrum.LatticeBasis{ByReciprocalSpace,D}`, which are aliased by `RealBasis{D}` and
+`ReciprocalBasis{D}` respectively.
+``
 """
-DataSpace(x) = DataSpace(typeof(basis(x)))
-DataSpace(::Type{T}) where T = DataSpace(fieldtype(T, :basis))
+DataSpace(::Type{<:LatticeBasis{S,D}}) where {S,D} = S{D}()
 
 #---Type promotion---------------------------------------------------------------------------------#
 
diff --git a/src/traits.jl b/src/traits.jl
index 456466bf..f8fa67b3 100644
--- a/src/traits.jl
+++ b/src/traits.jl
@@ -48,6 +48,26 @@ Trait for reciprocal space data in `D` dimensions.
 struct ByReciprocalSpace{D} <: BySpace{D}
 end
 
+"""
+    Electrum.DataSpace(x) -> DataSpace
+
+Returns a trait that determines whether a data set associated with a crystal is defined in real
+space (`ByRealSpace{D}()`), reciprocal space (`ByReciprocalSpace{D}()`), or by atomic positions
+(`ByAtom{D}`), where `D` is the number of dimensions.
+
+By default, `DataSpace(x)` will infer the appropriate trait from the lattice basis vectors
+included in `x`, assumed to be in a field named `basis`. The fallback definition is:
+
+    DataSpace(x) = DataSpace(typeof(x)) # basis(x) falls back to x.basis
+    DataSpace(::Type{T}) where T = DataSpace(fieldtype(T, :basis))
+
+For types `T` which use a `DataSpace` trait, but do not contain a set of lattice basis vectors which
+are `Electrum.LatticeBasis` objects stored in the `basis` field, it will be necessary to define
+`DataSpace(::Type{T})` explicitly for that type.
+"""
+DataSpace(x) = DataSpace(typeof(x))
+DataSpace(T::Type) = DataSpace(fieldtype(T, :basis))
+
 """
     Electrum.dimension(::DataSpace{D}) = D
     Electrum.dimension(::Type{<:DataSpace{D}}) = D